[rocm-libraries] ROCm/rocm-libraries#6073 (commit a538084)

[hipTensor] Add compiling options for unary ops
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

Add compiling options for unary ops and make hipTensor compiling fast

## Technical Details

This PR resolves an issue where the introduction of unary ops in
contraction caused hipTensor compile times to increase by 3–4x.
A new compile-time option, HIPTENSOR_INLINE_UNARY_OPS, has been added to
control this behavior. When set to ON, the switch_op method in the
HiptensorUnaryOp class is compiled inline, which was identified as the
root cause of the compile time regression. When set to OFF, switch_op is
compiled with noinline, restoring compile times to their baseline levels
prior to the addition of unary op support.
The default value of HIPTENSOR_INLINE_UNARY_OPS is OFF.

## Test Plan

run normal tests.
test compiling when HIPTENSOR_INLINE_UNARY_OPS is OFF or ON.

## Test Result

all tests passed.

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Author: Ryker0627

CMakeLists.txt

@@ -74,6 +74,7 @@ if(CMAKE_PROJECT_NAME STREQUAL "hiptensor")
     option(HIPTENSOR_DEFAULT_STRIDES_COL_MAJOR "Set hiptensor default strides to column major" ON)
     option(BUILD_OFFLOAD_COMPRESS "Build hiptensor with offload compression" ON)
     option(HIPTENSOR_CODE_COVERAGE "Build with code coverage flags (clang only)" OFF)
+    option(HIPTENSOR_INLINE_UNARY_OPS "Inline all unary ops for best runtime performance (slower compilation)" OFF)
 endif()
 
 # Setup output paths
@@ -157,6 +158,13 @@ else()
 endif()
 message(STATUS "HIPTENSOR_DEFAULT_STRIDES_COL_MAJOR=${HIPTENSOR_DEFAULT_STRIDES_COL_MAJOR}")
 
+if(HIPTENSOR_INLINE_UNARY_OPS)
+  add_compile_definitions(HIPTENSOR_INLINE_UNARY_OPS=1)
+else()
+  add_compile_definitions(HIPTENSOR_INLINE_UNARY_OPS=0)
+endif()
+message(STATUS "HIPTENSOR_INLINE_UNARY_OPS=${HIPTENSOR_INLINE_UNARY_OPS}")
+
 # Setup HIP
 find_package(hip REQUIRED)
 math(EXPR hip_VERSION_FLAT "(${hip_VERSION_MAJOR} * 1000 + ${hip_VERSION_MINOR}) * 100000 + ${hip_VERSION_PATCH}")

README.md

@@ -34,13 +34,14 @@ For more detailed information, please refer to the [hipTensor installation guide
 
 ### Project options
 
-| Option                              | Description                                             | Default Value                                           |
-|-------------------------------------|---------------------------------------------------------|---------------------------------------------------------|
-| GPU_TARGETS                         | Build code for specific GPU target(s)                   | gfx908;gfx90a;gfx942;gfx950;gfx11-generic;gfx12-generic |
-| HIPTENSOR_BUILD_TESTS               | Build the tests                                         | ON                                                      |
-| HIPTENSOR_BUILD_SAMPLES             | Build the samples                                       | ON                                                      |
-| HIPTENSOR_BUILD_COMPRESSED_DBG      | Enable compressed debug symbols                         | ON                                                      |
-| HIPTENSOR_DEFAULT_STRIDES_COL_MAJOR | Set the hipTensor default data layout to column major   | ON                                                      |
+| Option                              | Description                                                              | Default Value                                           |
+|-------------------------------------|--------------------------------------------------------------------------|---------------------------------------------------------|
+| GPU_TARGETS                         | Build code for specific GPU target(s)                                    | gfx908;gfx90a;gfx942;gfx950;gfx11-generic;gfx12-generic |
+| HIPTENSOR_BUILD_TESTS               | Build the tests                                                          | ON                                                      |
+| HIPTENSOR_BUILD_SAMPLES             | Build the samples                                                        | ON                                                      |
+| HIPTENSOR_BUILD_COMPRESSED_DBG      | Enable compressed debug symbols                                          | ON                                                      |
+| HIPTENSOR_DEFAULT_STRIDES_COL_MAJOR | Set the hipTensor default data layout to column major                    | ON                                                      |
+| HIPTENSOR_INLINE_UNARY_OPS          | Inline all unary ops for best runtime performance (slower compilation)   | OFF                                                     |
 
 ### Example configurations
 

docs/install/installation.rst

@@ -184,6 +184,9 @@ Here are the available options to build the hipTensor library, with or without c
     *   -   ``HIPTENSOR_DEFAULT_STRIDES_COL_MAJOR``
         -   Set the hipTensor default data layout to column major
         -   ``ON``
+    *   -   ``HIPTENSOR_INLINE_UNARY_OPS``
+        -   Inline all contraction unary ops for best runtime performance (slower compilation)
+        -   ``OFF``
 
 Here are some example project configurations:
 

library/src/include/hiptensor_element_wise_operation.hpp

@@ -36,6 +36,12 @@
 #include <ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp>
 #include <hiptensor/hiptensor_types.h>
 
+#if HIPTENSOR_INLINE_UNARY_OPS
+#define HIPTENSOR_UNARY_INLINE
+#else
+#define HIPTENSOR_UNARY_INLINE __attribute__((noinline))
+#endif
+
 namespace ck
 {
     namespace tensor_operation
@@ -179,7 +185,7 @@ namespace ck
                     = default;
 
                 template <typename T>
-                __host__ __device__ void switch_op(T& y, T const& x) const
+                __host__ __device__ HIPTENSOR_UNARY_INLINE void switch_op(T& y, T const& x) const
                 {
                     switch(op_type)
                     {