Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

A specialized Winograd Conv2d op #971

Draft
wants to merge 15 commits into
base: master
Choose a base branch
from
+1,417 −3
Draft

A specialized Winograd Conv2d op #971

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
5b4e448
added two winograd ops
Sep 26, 2024
68c251b
added more checking for precondition to use winograd; switch to im2co…
Sep 27, 2024
2ccc67d
added source code for winograd kernel
Sep 27, 2024
02a3cb1
winograd build ok
Sep 27, 2024
3d80466
test now passed; for some reason, ggml_conv_2d didn't output correct …
Sep 27, 2024
893ca79
remove debugging printouts
Sep 27, 2024
6afbf6e
added a FP16 FX kernel to deal with fp16 filter data; no need to use …
Sep 28, 2024
0491858
sync from sd.cpp
Sep 29, 2024
93c3da7
restore test-conv2d.cpp test
Sep 29, 2024
4e8e0d4
restore src/CMakeLists.txt
Sep 29, 2024
e0e94c4
change mask to unsigned int; add __restric__ to various pointers
Sep 29, 2024
c8700ca
fix indentation
Sep 29, 2024
c5d43a2
skip if already computed in a preprocess step
Oct 1, 2024
00ad37e
added winograd conv2d to backend op tests
Oct 1, 2024
4f93d67
add conv2d as a test case in test-backend-op
Oct 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -510,6 +510,8 @@ extern "C" {
GGML_OP_TIMESTEP_EMBEDDING,
GGML_OP_ARGSORT,
GGML_OP_LEAKY_RELU,
GGML_OP_WINOGRAD_STAGE0,
GGML_OP_WINOGRAD_STAGE1,

GGML_OP_FLASH_ATTN_EXT,
GGML_OP_FLASH_ATTN_BACK,
Expand Down Expand Up @@ -1696,6 +1698,21 @@ extern "C" {
struct ggml_tensor * a,
struct ggml_tensor * b,
int stride);

GGML_API struct ggml_tensor * ggml_winograd_stage0(
struct ggml_context * ctx,
struct ggml_tensor * a);

GGML_API struct ggml_tensor * ggml_winograd_stage1(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b);

GGML_API struct ggml_tensor * ggml_conv_2d_3x3(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b);


enum ggml_op_pool {
GGML_OP_POOL_MAX,
Expand Down
9 changes: 9 additions & 0 deletions src/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include "ggml-cuda/clamp.cuh"
#include "ggml-cuda/concat.cuh"
#include "ggml-cuda/conv-transpose-1d.cuh"
#include "ggml-cuda/conv-winograd.cuh"
#include "ggml-cuda/convert.cuh"
#include "ggml-cuda/cpy.cuh"
#include "ggml-cuda/cross-entropy-loss.cuh"
Expand Down Expand Up @@ -2331,6 +2332,12 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_OP_CONV_TRANSPOSE_1D:
ggml_cuda_op_conv_transpose_1d(ctx,dst);
break;
case GGML_OP_WINOGRAD_STAGE0:
bssrdf marked this conversation as resolved.
Show resolved Hide resolved
ggml_cuda_op_winograd_stage0(ctx, dst);
break;
case GGML_OP_WINOGRAD_STAGE1:
ggml_cuda_op_winograd_stage1(ctx, dst);
break;
case GGML_OP_POOL_2D:
ggml_cuda_op_pool2d(ctx, dst);
break;
Expand Down Expand Up @@ -2950,6 +2957,8 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
}
return false;
} break;
case GGML_OP_WINOGRAD_STAGE0:
case GGML_OP_WINOGRAD_STAGE1:
case GGML_OP_NONE:
case GGML_OP_RESHAPE:
case GGML_OP_VIEW:
Expand Down
851 changes: 851 additions & 0 deletions src/ggml-cuda/conv-winograd.cu
View file
Open in desktop

Large diffs are not rendered by default.

44 changes: 44 additions & 0 deletions src/ggml-cuda/conv-winograd.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
#include "common.cuh"


#define BC 8
#define BN 32
#define BK 64
#define TW 8
#define TH 16
#define BN_p 138

__constant__ int access_f_s[2][32];
__constant__ int access_s[2][32];
__constant__ int tileid[2][32];
Comment on lines +11 to +13
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens in the case of multiple GPUs? Is the constant memory duplicated across GPUs?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am pretty ignorant about multi-gpu. I guess they will be duplicated. I don't have a setup to test. Plus, this kernel only works for single GPU, I think.



// access_f_s
const int aux[2][32] = {
{0,0,1,1,2,2,3,3,4,4,5,5,6,6,
7,7,0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7},
{8,8,9,9,10,10,11,11,12,12,13,13,
14,14,15,15,8,8,9,9,10,10,11,11,12,12,
13,13,14,14,15,15}
};
// access_s
const int aux2[2][32] = {
{0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1,2,
3,2,3,2,3,2,3,2,3,2,3,2,3,2,3},
{4,5,4,5,4,5,4,5,4,5,4,
5,4,5,4,5,6,7,6,7,6,7,6,7,
6,7,6,7,6,7,6,7}
};

const int tid[2][32] = {
{0,1,4,5,8,9,12,13,16,17,20,21,24,25,28,29,
0,1,4,5,8,9,12,13,16,17,20,21,24,25,28,29},
{2,3,6,7,10,11,14,15,18,19,22,23,26,27,30,31,
2,3,6,7,10,11,14,15,18,19,22,23,26,27,30,31}
};



void ggml_cuda_op_winograd_stage0(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
void ggml_cuda_op_winograd_stage1(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

108 changes: 106 additions & 2 deletions src/ggml.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2995,6 +2995,8 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"TIMESTEP_EMBEDDING",
"ARGSORT",
"LEAKY_RELU",
"WINOGRAD_STAGE0",
"WINOGRAD_STAGE1",

"FLASH_ATTN_EXT",
"FLASH_ATTN_BACK",
Expand Down Expand Up @@ -3024,7 +3026,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"OPT_STEP_ADAMW",
};

static_assert(GGML_OP_COUNT == 80, "GGML_OP_COUNT != 80");
static_assert(GGML_OP_COUNT == 82, "GGML_OP_COUNT != 82");

static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"none",
Expand Down Expand Up @@ -3089,6 +3091,8 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"timestep_embedding(timesteps, dim, max_period)",
"argsort(x)",
"leaky_relu(x)",
"winograd_stage0(x)",
"winograd_stage1(x)",

"flash_attn_ext(x)",
"flash_attn_back(x)",
Expand Down Expand Up @@ -3118,7 +3122,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"adamw(x)",
};

static_assert(GGML_OP_COUNT == 80, "GGML_OP_COUNT != 80");
static_assert(GGML_OP_COUNT == 82, "GGML_OP_COUNT != 82");

static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");

Expand Down Expand Up @@ -7166,6 +7170,73 @@ struct ggml_tensor * ggml_conv_transpose_2d_p0(
return result;
}


// ggml_winograd

// a: [OC，IC, 3, 3]
// result: [OC, IC, 16]
struct ggml_tensor * ggml_winograd_stage0(
struct ggml_context * ctx,
struct ggml_tensor * a) {
bool is_node = false;

if (a->grad) {
is_node = true;
}
Comment on lines +7181 to +7185
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If #966 is merged first this will need to be removed (should be very straightforward).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Look forward to it...


struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, a->ne[3], 4, 4, a->ne[2]);

result->op = GGML_OP_WINOGRAD_STAGE0;
result->grad = is_node ? ggml_dup_tensor(ctx, result) : NULL;
result->src[0] = a;

return result;
}

// ggml_winograd
// a: [OC, IC, 4, 4]
// b: [1, IC, IH, IW]
// result: [N, OC, OH, OW]
struct ggml_tensor * ggml_winograd_stage1(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b) {
bool is_node = false;
if (a->grad) {
is_node = true;
}

int OW = b->ne[0];
int OH = b->ne[1];
struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, OW, OH, a->ne[0] /* OC */, 1);

result->op = GGML_OP_WINOGRAD_STAGE1;
result->grad = is_node ? ggml_dup_tensor(ctx, result) : NULL;
result->src[0] = a;
result->src[1] = b;

return result;
}

struct ggml_tensor * ggml_conv_2d_3x3(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b){
GGML_ASSERT(a->ne[0] == 3 && a->ne[1] == 3); // kernel should be 3x3
GGML_ASSERT(b->ne[3] == 1); // only works for 1 input image
GGML_ASSERT(b->ne[2] == a->ne[2]); // number of channels must match
if(a->ne[3] % 64 != 0 || a->ne[2] % 8 != 0) // only works for the number of filters is a multiple of 64
return ggml_conv_2d(ctx, a, b, 1, 1, 1, 1, 1, 1); // and the number of channels is a multiple of 8

// struct ggml_tensor* ra = ggml_cont(ctx, ggml_permute(ctx, a, 1, 2, 3, 0)); // [N, OC, OH, OW]
struct ggml_tensor* W = ggml_winograd_stage0(ctx, a);
struct ggml_tensor * result = ggml_winograd_stage1(ctx, W, b);

return result;

}


// ggml_pool_*

static int64_t ggml_calc_pool_output_size(int64_t ins, int ks, int s, float p) {
Expand Down Expand Up @@ -15124,6 +15195,23 @@ static void ggml_compute_forward_conv_transpose_1d(
}
}


static void ggml_compute_forward_winograd_stage0(
const struct ggml_compute_params * params,
struct ggml_tensor * dst) {

GGML_ASSERT(false && " CPU backend not implemented!");
return;
}

static void ggml_compute_forward_winograd_stage1(
const struct ggml_compute_params * params,
struct ggml_tensor * dst) {

GGML_ASSERT(false && " CPU backend not implemented!");
return;
}
Comment on lines +15199 to +15213
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If at all possible a CPU implementation should always be done since it serves both as a fallback and as a reference implementation to test other backends against.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A CPU backend should be done, but I am not sure the benefit of it compared to the current im2col+gemm version.


// ggml_compute_forward_im2col_f32
// src0: kernel [OC, IC, KH, KW]
// src1: image [N, IC, IH, IW]
Expand Down Expand Up @@ -17820,6 +17908,14 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
{
ggml_compute_forward_conv_transpose_1d(params, tensor);
} break;
case GGML_OP_WINOGRAD_STAGE0:
{
ggml_compute_forward_winograd_stage0(params, tensor);
} break;
case GGML_OP_WINOGRAD_STAGE1:
{
ggml_compute_forward_winograd_stage1(params, tensor);
} break;
case GGML_OP_IM2COL:
{
ggml_compute_forward_im2col(params, tensor);
Expand Down Expand Up @@ -18893,6 +18989,14 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
{
GGML_ABORT("fatal error"); // TODO: not implemented
}
case GGML_OP_WINOGRAD_STAGE0:
{
GGML_ABORT("fatal error"); // TODO: not implemented
}
case GGML_OP_WINOGRAD_STAGE1:
{
GGML_ABORT("fatal error"); // TODO: not implemented
}
case GGML_OP_POOL_1D:
{
GGML_ABORT("fatal error"); // TODO: not implemented
Expand Down
9 changes: 9 additions & 0 deletions tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,15 @@ target_link_libraries(${TEST_TARGET} PRIVATE ggml)
add_test(NAME ${TEST_TARGET} COMMAND $<TARGET_FILE:${TEST_TARGET}>)
set_property(TEST ${TEST_TARGET} PROPERTY ENVIRONMENT "LLVM_PROFILE_FILE=${TEST_TARGET}.profraw")

#
# test-conv2d-wino

set(TEST_TARGET test-conv2d-winograd)
add_executable(${TEST_TARGET} ${TEST_TARGET}.cpp)
target_link_libraries(${TEST_TARGET} PRIVATE ggml)
add_test(NAME ${TEST_TARGET} COMMAND $<TARGET_FILE:${TEST_TARGET}>)
set_property(TEST ${TEST_TARGET} PROPERTY ENVIRONMENT "LLVM_PROFILE_FILE=${TEST_TARGET}.profraw")


#
# test-mul-mat
Expand Down
Loading
Loading