mirror of
https://github.com/ml-explore/mlx.git
synced 2025-12-16 01:49:05 +08:00
17 lines
360 B
C++
17 lines
360 B
C++
// Copyright © 2025 Apple Inc.
|
|
|
|
#pragma once
|
|
|
|
#include <hip/hip_runtime.h>
|
|
|
|
namespace mlx::core::rocm {
|
|
|
|
template <typename T>
|
|
__global__ void arange_kernel(T* out, T start, T step, size_t size) {
|
|
size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (tid < size) {
|
|
out[tid] = start + static_cast<T>(tid) * step;
|
|
}
|
|
}
|
|
|
|
} // namespace mlx::core::rocm
|