fix 115

Grid_sample.mlu

@@ -0,0 +1,132 @@
+#include <bang.h>
+#include <torch/extension.h>
+#include <cnrt.h>
+
+/* ============================================================================
+ * 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<<<dim, ktype, queue>>>(
+        reinterpret_cast<const half *>(input.data_ptr<at::Half>()),
+        reinterpret_cast<const half *>(grid.data_ptr<at::Half>()),
+        reinterpret_cast<half *>(output.data_ptr<at::Half>()),
+        N,
+        C,
+        H,
+        W,
+        out_H,
+        out_W,
+        N * out_H * out_W);
+
+    return output;
+}

Unfold.mlu

@@ -0,0 +1,165 @@
+#include <bang.h>
+#include <torch/extension.h>
+#include <cnrt.h>
+
+/* ============================================================================
+ * 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<<<dim, ktype, queue>>>(
+            reinterpret_cast<const half *>(x.data_ptr<at::Half>()),
+            reinterpret_cast<half *>(output.data_ptr<at::Half>()),
+            N,
+            C,
+            H,
+            W,
+            K,
+            H_out,
+            W_out,
+            total_rows);
+    } else {
+        unfold_kernel<<<dim, ktype, queue>>>(
+            reinterpret_cast<const half *>(x.data_ptr<at::Half>()),
+            reinterpret_cast<half *>(output.data_ptr<at::Half>()),
+            N,
+            C,
+            H,
+            W,
+            K,
+            stride,
+            padding,
+            H_out,
+            W_out,
+            total);
+    }
+
+    return output;
+}

config

@@ -1 +1,2 @@
-075
+115
+116