diff --git a/Grid_sample.mlu b/Grid_sample.mlu new file mode 100644 index 0000000..f88cf7c --- /dev/null +++ b/Grid_sample.mlu @@ -0,0 +1,132 @@ +#include +#include +#include + +/* ============================================================================ + * Grid Sample (bilinear, zeros padding, align_corners=True) + * + * input: [N, C, H, W] + * grid: [N, out_H, out_W, 2] + * output: [N, C, out_H, out_W] + * ============================================================================ + */ +__mlu_entry__ void grid_sample_bilinear_kernel( + const half *input, + const half *grid, + half *output, + int N, + int C, + int H, + int W, + int out_H, + int out_W, + int total_grid) +{ + int task_id = (int)taskId; + int task_num = (int)taskDim; + + int in_hw = H * W; + int out_hw = out_H * out_W; + + for (int pos = task_id; pos < total_grid; pos += task_num) { + int n = pos / out_hw; + int out_index = pos - n * out_hw; + int oh = out_index / out_W; + int ow = out_index - oh * out_W; + + int grid_base = ((n * out_H + oh) * out_W + ow) * 2; + float gx = (float)grid[grid_base]; + float gy = (float)grid[grid_base + 1]; + + float ix = (gx + 1.0f) * (float)(W - 1) * 0.5f; + float iy = (gy + 1.0f) * (float)(H - 1) * 0.5f; + + int ix0 = (int)ix; + int iy0 = (int)iy; + if ((float)ix0 > ix) ix0 -= 1; + if ((float)iy0 > iy) iy0 -= 1; + + int ix1 = ix0 + 1; + int iy1 = iy0 + 1; + + float wx1 = ix - (float)ix0; + float wy1 = iy - (float)iy0; + float wx0 = 1.0f - wx1; + float wy0 = 1.0f - wy1; + + int valid00 = (iy0 >= 0 && iy0 < H && ix0 >= 0 && ix0 < W); + int valid01 = (iy0 >= 0 && iy0 < H && ix1 >= 0 && ix1 < W); + int valid10 = (iy1 >= 0 && iy1 < H && ix0 >= 0 && ix0 < W); + int valid11 = (iy1 >= 0 && iy1 < H && ix1 >= 0 && ix1 < W); + + float w00 = wy0 * wx0; + float w01 = wy0 * wx1; + float w10 = wy1 * wx0; + float w11 = wy1 * wx1; + + for (int c = 0; c < C; ++c) { + const half *in_ch = input + (n * C + c) * in_hw; + float acc = 0.0f; + + if (valid00) { + acc += (float)in_ch[iy0 * W + ix0] * w00; + } + if (valid01) { + acc += (float)in_ch[iy0 * W + ix1] * w01; + } + if (valid10) { + acc += (float)in_ch[iy1 * W + ix0] * w10; + } + if (valid11) { + acc += (float)in_ch[iy1 * W + ix1] * w11; + } + + output[(n * C + c) * out_hw + out_index] = (half)acc; + } + } +} + + +torch::Tensor bang_func(torch::Tensor input, torch::Tensor grid) +{ + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(grid.is_contiguous(), "grid must be contiguous"); + TORCH_CHECK(input.scalar_type() == torch::kFloat16, "input must be float16"); + TORCH_CHECK(grid.scalar_type() == torch::kFloat16, "grid must be float16"); + TORCH_CHECK(input.dim() == 4, "input must be 4D: [N, C, H, W]"); + TORCH_CHECK(grid.dim() == 4, "grid must be 4D: [N, out_H, out_W, 2]"); + TORCH_CHECK(grid.size(3) == 2, "grid last dimension must be 2"); + TORCH_CHECK(input.size(0) == grid.size(0), "input and grid batch size must match"); + + int N = (int)input.size(0); + int C = (int)input.size(1); + int H = (int)input.size(2); + int W = (int)input.size(3); + int out_H = (int)grid.size(1); + int out_W = (int)grid.size(2); + + TORCH_CHECK(N > 0 && C > 0 && H > 0 && W > 0, + "input dimensions must be greater than 0"); + TORCH_CHECK(out_H > 0 && out_W > 0, + "grid output dimensions must be greater than 0"); + + auto output = torch::empty({N, C, out_H, out_W}, input.options()); + + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {16, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeBlock; + + grid_sample_bilinear_kernel<<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(grid.data_ptr()), + reinterpret_cast(output.data_ptr()), + N, + C, + H, + W, + out_H, + out_W, + N * out_H * out_W); + + return output; +} diff --git a/Unfold.mlu b/Unfold.mlu new file mode 100644 index 0000000..16e2f29 --- /dev/null +++ b/Unfold.mlu @@ -0,0 +1,165 @@ +#include +#include +#include + +/* ============================================================================ + * Unfold / im2col + * + * input: [N, C, H, W] + * output: [N, C * K * K, H_out * W_out] + * + * Layout matches torch.nn.Unfold(kernel_size=K, stride=stride, padding=padding). + * ============================================================================ + */ +__mlu_entry__ void unfold_kernel( + const half *input, + half *output, + int N, + int C, + int H, + int W, + int K, + int stride, + int padding, + int H_out, + int W_out, + int total) +{ + int task_id = (int)taskId; + int task_num = (int)taskDim; + + int columns = H_out * W_out; + int kernel_area = K * K; + int rows = C * kernel_area; + int per_batch = rows * columns; + + for (int index = task_id; index < total; index += task_num) { + int n = index / per_batch; + int inner = index - n * per_batch; + int row = inner / columns; + int col = inner - row * columns; + + int c = row / kernel_area; + int k_rem = row - c * kernel_area; + int kh = k_rem / K; + int kw = k_rem - kh * K; + + int oh = col / W_out; + int ow = col - oh * W_out; + + int ih = oh * stride + kh - padding; + int iw = ow * stride + kw - padding; + + if (ih >= 0 && ih < H && iw >= 0 && iw < W) { + int input_index = ((n * C + c) * H + ih) * W + iw; + output[index] = input[input_index]; + } else { + output[index] = (half)0.0f; + } + } +} + + +__mlu_entry__ void unfold_stride1_pad0_kernel( + const half *input, + half *output, + int N, + int C, + int H, + int W, + int K, + int H_out, + int W_out, + int total_rows) +{ + int task_id = (int)taskId; + int task_num = (int)taskDim; + + int kernel_area = K * K; + int rows_per_batch = C * kernel_area; + int out_columns = H_out * W_out; + int items_per_batch = rows_per_batch * H_out; + + for (int item = task_id; item < total_rows; item += task_num) { + int n = item / items_per_batch; + int inner = item - n * items_per_batch; + int row = inner / H_out; + int oh = inner - row * H_out; + + int c = row / kernel_area; + int k_rem = row - c * kernel_area; + int kh = k_rem / K; + int kw = k_rem - kh * K; + + const half *in_ptr = input + ((n * C + c) * H + (oh + kh)) * W + kw; + half *out_ptr = output + (n * rows_per_batch + row) * out_columns + oh * W_out; + + for (int ow = 0; ow < W_out; ++ow) { + out_ptr[ow] = in_ptr[ow]; + } + } +} + + +torch::Tensor bang_func(torch::Tensor x, int kernel_size, int stride, int padding) +{ + TORCH_CHECK(x.is_contiguous(), "x must be contiguous"); + TORCH_CHECK(x.scalar_type() == torch::kFloat16, "x must be float16"); + TORCH_CHECK(x.dim() == 4, "x must be 4D: [N, C, H, W]"); + TORCH_CHECK(kernel_size > 0, "kernel_size must be greater than 0"); + TORCH_CHECK(stride > 0, "stride must be greater than 0"); + TORCH_CHECK(padding >= 0, "padding must be non-negative"); + + int N = (int)x.size(0); + int C = (int)x.size(1); + int H = (int)x.size(2); + int W = (int)x.size(3); + int K = kernel_size; + + int H_out = (H + 2 * padding - K) / stride + 1; + int W_out = (W + 2 * padding - K) / stride + 1; + + TORCH_CHECK(H_out > 0 && W_out > 0, + "Invalid output size: H_out=", H_out, ", W_out=", W_out); + + int rows = C * K * K; + int columns = H_out * W_out; + int total = N * rows * columns; + + auto output = torch::empty({N, rows, columns}, x.options()); + + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {16, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeBlock; + + if (padding == 0 && stride == 1) { + int total_rows = N * rows * H_out; + unfold_stride1_pad0_kernel<<>>( + reinterpret_cast(x.data_ptr()), + reinterpret_cast(output.data_ptr()), + N, + C, + H, + W, + K, + H_out, + W_out, + total_rows); + } else { + unfold_kernel<<>>( + reinterpret_cast(x.data_ptr()), + reinterpret_cast(output.data_ptr()), + N, + C, + H, + W, + K, + stride, + padding, + H_out, + W_out, + total); + } + + return output; +} diff --git a/config b/config index 920a6ea..7a6719f 100644 --- a/config +++ b/config @@ -1 +1,2 @@ -075 +115 +116