optimize softmax gpu implementation.

src/targets/gpu/device/softmax.cpp

@@ -30,45 +30,98 @@ argument softmax(hipStream_t stream,
             hip_tensor_descriptor<n_dim> desc_batch(batch_shape);
             hip_tensor_descriptor<n_dim> desc_data(output_shape);
 
-            // each thread is for one item in the batch
-            gs_launch(stream, batch_shape.elements())([=](auto i) {
-                auto batch_idx = desc_batch.multi(i);
+            // use one block for items in one batch.
+            const size_t block_size = 1024;
+            launch(
+                stream, batch_shape.elements() * block_size, block_size)([=](auto idx) __device__ {
+                size_t thr_idx = idx.local;
+                size_t blk_idx = idx.group;
+                using type = device_type<std::remove_cv_t<typename decltype(output)::value_type>>;
+
+                // all data can be loaded to the lds once, so all operations are
+                // done in lds
+                MIGRAPHX_DEVICE_SHARED type lds_data[block_size + 2];
+                auto batch_idx = desc_batch.multi(blk_idx);
                 auto data_idx  = batch_idx;
-                // get max
-                auto batch_max = input_ptr[desc_data.linear(batch_idx)];
-                for(std::size_t j = 1; j < n_dims; ++j)
+                // load data to lds and compute the batch max
+                size_t item_num      = n_dims;
+                lds_data[block_size] = input_ptr[0];
+                for(size_t i = thr_idx; i < n_dims; i += block_size)
                 {
-                    data_idx[axis] = j;
-                    batch_max      = std::max(to_hip_type(batch_max),
-                                         to_hip_type(input_ptr[desc_data.linear(data_idx)]));
-                }
+                    data_idx[axis] = i;
+                    lds_data[i]    = input_ptr[desc_data.linear(data_idx)];
 
-                for(std::size_t j = 0; j < n_dims; ++j)
-                {
-                    data_idx[axis]  = j;
-                    auto idx        = desc_data.linear(data_idx);
-                    output_ptr[idx] = input_ptr[idx] - batch_max;
-                }
+                    __syncthreads();
 
-                for(std::size_t j = 0; j < n_dims; ++j)
-                {
-                    data_idx[axis]  = j;
-                    auto idx        = desc_data.linear(data_idx);
-                    output_ptr[idx] = exp(to_hip_type(output_ptr[idx]));
+                    auto size   = (item_num > block_size) ? block_size : item_num;
+                    auto stride = (size + 1) / 2;
+                    while(true)
+                    {
+                        if(thr_idx + stride < size)
+                        {
+                            lds_data[thr_idx] = ::max(to_hip_type(lds_data[thr_idx]),
+                                                      to_hip_type(lds_data[thr_idx + stride]));
+                        }
+                        __syncthreads();
+                        size   = stride;
+                        stride = (stride + 1) / 2;
+
+                        if(size == 1)
+                            break;
+                    }
+
+                    if(thr_idx == 0)
+                    {
+                        lds_data[block_size] = (lds_data[0] < lds_data[block_size])
+                                                   ? lds_data[block_size]
+                                                   : lds_data[0];
+                    }
+                    __syncthreads();
+
+                    item_num -= block_size;
                 }
 
-                auto batch_sum = output_ptr[desc_data.linear(batch_idx)];
-                for(std::size_t j = 1; j < n_dims; ++j)
+                const size_t block_size1 = block_size + 1;
+                lds_data[block_size1]    = 0;
+                item_num                 = n_dims;
+                for(size_t i = thr_idx; i < n_dims; i += block_size)
                 {
-                    data_idx[axis] = j;
-                    batch_sum += output_ptr[desc_data.linear(data_idx)];
+                    data_idx[axis] = i;
+                    lds_data[i]    = input_ptr[desc_data.linear(data_idx)] - lds_data[block_size];
+                    lds_data[i]    = ::exp(to_hip_type(lds_data[i]));
+
+                    __syncthreads();
+
+                    auto size   = (item_num > block_size) ? block_size : item_num;
+                    auto stride = (size + 1) / 2;
+                    while(true)
+                    {
+                        if(thr_idx + stride < size)
+                        {
+                            lds_data[thr_idx] += lds_data[thr_idx + stride];
+                        }
+                        __syncthreads();
+                        size   = stride;
+                        stride = (stride + 1) / 2;
+                        if(size == 1)
+                            break;
+                    }
+
+                    if(thr_idx == 0)
+                    {
+                        lds_data[block_size1] += lds_data[0];
+                    }
+                    __syncthreads();
+
+                    item_num -= block_size;
                 }
 
-                for(std::size_t j = 0; j < n_dims; ++j)
+                for(size_t i = thr_idx; i < n_dims; i += block_size)
                 {
-                    data_idx[axis]  = j;
-                    auto idx        = desc_data.linear(data_idx);
-                    output_ptr[idx] = output_ptr[idx] / batch_sum;
+                    data_idx[axis]    = i;
+                    size_t index      = desc_data.linear(data_idx);
+                    auto val          = input_ptr[index] - lds_data[block_size];
+                    output_ptr[index] = ::exp(to_hip_type(val)) / lds_data[block_size1];
                 }
             });
         });