[ROCM] fix concat and split

[ronny1996]

PR types

Bug fixes

PR changes

Others

Description

fix concat and split in dtk23.04

#include <hip/hip_runtime.h>

#include <iostream>
#include <cstdio>

struct __attribute__((aligned(256))) Foo {
    Foo(int numel) : numel_(numel) {}

    __device__ int n() const {
        return numel_;
    }

    private:
    int numel_{0};
};

__global__ void add_kernel(const float* in_x, const float* in_y, float* out_z, Foo numel) {
    auto tid = threadIdx.x + blockIdx.x * blockDim.x;
    printf("(%d,%d,%d),(%d,%d,%d) tid=%d\n", 
    static_cast<int>(blockIdx.x), 
    static_cast<int>(blockIdx.y), 
    static_cast<int>(blockIdx.z),
    static_cast<int>(threadIdx.x), 
    static_cast<int>(threadIdx.y), 
    static_cast<int>(threadIdx.z), 
    tid);
    for (auto i = tid; i < numel.n(); i += blockDim.x * gridDim.x) {
        out_z[i] = in_x[i] + in_y[i];
    }
}


int main() {
    const size_t numel = 1024;
    const size_t bytesize = numel * sizeof(float);
    float* host_x, *host_y, *host_z;
    float* device_x, *device_y, *device_z;

    host_x = static_cast<float*>(malloc(bytesize));
    host_y = static_cast<float*>(malloc(bytesize));
    host_z = static_cast<float*>(malloc(bytesize));
    hipMalloc(&device_x, bytesize);
    hipMalloc(&device_y, bytesize);
    hipMalloc(&device_z, bytesize);
    for (auto i = 0; i < numel; ++i) {
        host_x[i] = 1;
        host_y[i] = 2;
    }

    hipMemcpy(device_x, host_x, bytesize, hipMemcpyHostToDevice);
    hipMemcpy(device_y, host_y, bytesize, hipMemcpyHostToDevice);
    Foo foo(numel);
    add_kernel<<<16, 512>>>(device_x, device_y, device_z, foo);
    hipMemcpy(host_z, device_z, bytesize, hipMemcpyHostToDevice);

    for (auto i = 0; i < numel; ++i) {
        std::cout << host_z[i] << ", ";
    }
    std::cout << std::endl;
    free(host_x);
    free(host_y);
    free(host_z);
    hipFree(device_x);
    hipFree(device_y);
    hipFree(device_z);
    return 0;
}

dtk23.04 中 __attribute__((aligned(256))) 会导致 segmentation fault

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[qili93]

LGTM