Skip to content

Commit 078df41

Browse files
dzzz2001claude
dzzz2001
and
claude
authored
perf(gint): shape-exact bucketing + tile ladder + wide-LDS vbatched GEMM (#7395)
* perf(gint): shape-exact bucketing + tile ladder + wide-LDS vbatched GEMM Optimize the GPU gint batched-GEMM path (gemm_{nn,tn}_vbatch, driven from phi_mul_phi / phi_mul_dm) for FP64 on V100/A100-class GPUs. - phi_operator_gpu: replace the single max-shape vbatch launch with shape-exact bucketing. Atom pairs are grouped by (nw1, nw2) via a dense NW_MAX*NW_MAX counting-sort table, pre-enumerated once per batch in set_bgrid_batch, so each bucket hands the kernel a scalar (m, n, k) and the tile ladder picks the tightest tile per shape -- no cross-species tile waste, no over-launched blocks. A guard aborts if any atom nw >= NW_MAX. - dgemm_vbatch: scalar (m, n, k) dispatch (drops the per-batchid M/N/K device arrays) feeding a 4x2 (NN) / 4x4 (TN) BLK_{M,N} ladder over {8,16,32,48}. - gemm_{nn,tn}_vbatch: K-inner shared-memory layout + wide (double2/float4) LDS inner loop -- one 16-byte LDS feeds VK FMAs per (m,n); PAD keeps the shmem stride 16-byte aligned and warp access bank-conflict-free. C accumulators stay double regardless of input type T, preserving the mixed-precision fp64-accumulator fix (#7368); the phi_operator kernel optimizations from #7366 (WantPhi dispatch, single-warp reduce) are retained. FP64 15-case GPU benchmark: end-to-end ~1.05x (A800) / ~1.04x (V100), with cal_gint_vl up to ~1.5x and cal_gint_rho up to ~1.65x; energies and pressures match develop to ~1e-10 on every case. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(gint): derive shape-bucket stride from ucell.nwmax, drop hardcoded NW_MAX The (nw1, nw2) shape-bucketing in phi_mul_phi / phi_mul_dm flattened pairs into a dense table key via `nw1 * NW_MAX + nw2`, with NW_MAX a hardcoded 64. That was both a magic number and an artificial ceiling: a basis with nw > 64 would abort(), and 64 was only a guess at the real max. The true upper bound is already known to the code as ucell.nwmax (max orbital count over all atom types), exposed via gint_gpu_vars_->nwmax. Use it: set nw_stride_ = nwmax + 1 once in the ctor so the bucket table is sized exactly to the basis -- no cap to maintain. A runtime stride can't index std::array<int, NW_MAX*NW_MAX>, so the three counting-sort tables (counts / base / cursor) move to mutable std::vector members allocated once and re-zeroed per call. For typical nwmax~25 that's ~676 ints vs the old fixed 4096, so the hot path zeroes less and never reallocates. The set_bgrid_batch() abort guard becomes a structurally-unreachable assert, since nwmax is by definition the largest nw. Drop now-unused includes (<array>, <cstdio>, <cstdlib>); add <cassert>. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor(gint): clarify GEMM kernel comments, hoist shape-bucket struct Follow-up cleanup on the shape-exact vbatched GEMM path. No behavior change. - gemm_{nn,tn}_vbatch, dgemm_vbatch, gint_helper: rewrite the kernel comments to describe the actual mechanism (K-inner shared-memory layout, wide vector loads feeding VK FMAs per load, the tile ladder, fp64 cross-item accumulation) and drop the internal "V1/V3/Phase" development shorthand that carried no meaning outside the original work log. - phi_operator_gpu: the local `Bucket` struct was declared identically inside both phi_mul_phi and phi_mul_dm. Hoist it to a named GemmShapeBucket type and reuse a single buckets_ member vector (cleared, not reallocated) across both, reserved once in the ctor -- one less per-call heap allocation on the hot path. - phi_operator_gpu: pair_scratch_offset_ is fully overwritten in Pass 1 before Pass 2 reads it, so resize() it instead of assign(..., -1); the -1 sentinel was never observed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 69a663f commit 078df41

7 files changed

Lines changed: 665 additions & 323 deletions

File tree

‎source/source_lcao/module_gint/kernel/dgemm_vbatch.cu‎

Lines changed: 108 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,60 +3,148 @@
33
#include "dgemm_vbatch.h"
44
#include "source_base/module_device/device.h"
55

6+
// Tile ladder
7+
// -----------
8+
// The caller splits each batch into buckets of identical (m, n, k) and calls
9+
// in once per bucket. The dispatchers below pick, for each bucket, the kernel
10+
// instantiation whose (BLK_M, BLK_N) tile is the smallest rung that still
11+
// covers the bucket's output shape, so boundary blocks don't spend most of
12+
// their work on masked-off padding.
13+
//
14+
// Each thread owns a THR_M x THR_N register accumulator tile, i.e. it computes
15+
// THR = THR_M * THR_N = (BLK_M / DIM_X) * (BLK_N / DIM_Y)
16+
// output elements. We aim to keep THR in roughly [16, 36]: below that the inner
17+
// FMAs don't amortize the shared-memory traffic and there's too little ILP;
18+
// above it register pressure starts cutting occupancy. The "(in band)" /
19+
// "(under)" notes on each case below mark where that rung lands.
20+
621
template<typename T>
722
void gemm_nn_vbatch(
8-
int max_m, int max_n, int max_k,
9-
const int* m_d, const int* n_d, const int* k_d,
23+
int m, int n, int k,
1024
const T* const* A_array_d, const int* lda_d,
1125
const T* const* B_array_d, const int* ldb_d,
1226
double** C_array_d, const int* ldc_d,
1327
int batchCount, cudaStream_t stream,
1428
const T* alpha)
1529
{
16-
vbatched_gemm_nn_impl<T, 8, 4, 16, 16, 8, 8, 4, 8, 4>
17-
(max_m, max_n, m_d, n_d, k_d,
18-
A_array_d, lda_d,
19-
B_array_d, ldb_d,
20-
C_array_d, ldc_d,
21-
batchCount, stream, alpha);
30+
// 4 (nw2 bracket) x 2 (bxyz bracket) = 8 instantiations.
31+
//
32+
// Mapping into the impl's parameter list is:
33+
// <T, DIM_X, DIM_Y, BLK_M, BLK_N, BLK_K=16,
34+
// DIM_XA=DIM_X, DIM_YA=DIM_Y, DIM_XB=DIM_X, DIM_YB=DIM_Y>
35+
// which satisfies the kernel's tile-divisibility asserts because every
36+
// (BLK_M, BLK_N, BLK_K=16) chosen below is a multiple of the matching
37+
// (DIM_X, DIM_Y) pair.
38+
#define NN_DISPATCH(DX, DY, BM, BN) \
39+
vbatched_gemm_nn_impl<T, DX, DY, BM, BN, 16, DX, DY, DX, DY>( \
40+
m, n, k, \
41+
A_array_d, lda_d, B_array_d, ldb_d, \
42+
C_array_d, ldc_d, batchCount, stream, alpha)
43+
44+
// BLK_M bracket -- smallest tile in {8,16,32,48} covering nw2.
45+
const int blk_m_tag = (n <= 8) ? 0
46+
: (n <= 16) ? 1
47+
: (n <= 32) ? 2
48+
: 3;
49+
50+
// BLK_N bracket -- tiles the bxyz (mesh-grid) axis. Use 32 when bxyz<=32 so
51+
// a partial final block-row isn't mostly masked padding (e.g. bxyz=27 in a
52+
// 64-row tile leaves ~58% of the rows idle); use 64 above that, where the
53+
// larger tile gives better shared-memory reuse.
54+
const int blk_n_tag = (m <= 32) ? 0 : 1;
55+
56+
switch (blk_m_tag * 2 + blk_n_tag)
57+
{
58+
// BLK_M=8 (nw2 <=8 ). DIM=4x8 -> THR_M=2.
59+
case 0: NN_DISPATCH( 4, 8, 8, 32); break; // THR=2*4=8 (under)
60+
case 1: NN_DISPATCH( 4, 8, 8, 64); break; // THR=2*8=16 (in band)
61+
// BLK_M=16 (nw2<=16). DIM=4x8 -> THR_M=4.
62+
case 2: NN_DISPATCH( 4, 8, 16, 32); break; // THR=4*4=16 (in band)
63+
case 3: NN_DISPATCH( 4, 8, 16, 64); break; // THR=4*8=32 (in band)
64+
// BLK_M=32 (nw2<=32). DIM=8x8 -> THR_M=4.
65+
case 4: NN_DISPATCH( 8, 8, 32, 32); break; // THR=4*4=16 (in band)
66+
case 5: NN_DISPATCH( 8, 8, 32, 64); break; // THR=4*8=32 (in band)
67+
// BLK_M=48 (nw2<=48). DIM=16x8 -> THR_M=3 (cap at 3 to keep
68+
// register pressure room for the BLK_N=64 sibling).
69+
case 6: NN_DISPATCH(16, 8, 48, 32); break; // THR=3*4=12 (just under)
70+
case 7: NN_DISPATCH(16, 8, 48, 64); break; // THR=3*8=24 (in band)
71+
}
2272

73+
#undef NN_DISPATCH
2374
}
2475

2576
template<typename T>
2677
void gemm_tn_vbatch(
27-
int max_m, int max_n, int max_k,
28-
const int* m_d, const int* n_d, const int* k_d,
78+
int m, int n, int k,
2979
const T* const* A_array_d, const int* lda_d,
3080
const T* const* B_array_d, const int* ldb_d,
3181
double** C_array_d, const int* ldc_d,
3282
int batchCount, cudaStream_t stream,
3383
const T* alpha)
3484
{
35-
vbatched_gemm_tn_impl<T, 8,4,16,16,4,8,4,8,4>
36-
(max_m, max_n, m_d, n_d, k_d,
37-
A_array_d, lda_d,
38-
B_array_d, ldb_d,
39-
C_array_d, ldc_d,
40-
batchCount, stream, alpha);
85+
// 4 (nw2 bracket) x 4 (nw1 bracket) = 16 instantiations.
86+
//
87+
// Both output axes here are the small nw axis, so we use the same
88+
// {8,16,32,48} ladder on both. BLK_K = 32 (the bxyz axis -- large).
89+
#define TN_DISPATCH(DX, DY, BM, BN) \
90+
vbatched_gemm_tn_impl<T, DX, DY, BM, BN, 32, DX, DY, DX, DY>( \
91+
m, n, k, \
92+
A_array_d, lda_d, B_array_d, ldb_d, \
93+
C_array_d, ldc_d, batchCount, stream, alpha)
94+
95+
auto bracket = [](int x) {
96+
return (x <= 8) ? 0
97+
: (x <= 16) ? 1
98+
: (x <= 32) ? 2
99+
: 3;
100+
};
101+
const int blk_m_tag = bracket(n); // BLK_M <- nw2
102+
const int blk_n_tag = bracket(m); // BLK_N <- nw1
103+
104+
switch (blk_m_tag * 4 + blk_n_tag)
105+
{
106+
// BLK_M=8 rungs (nw2<=8). DIM_X=4, THR_M=2.
107+
case 0: TN_DISPATCH(4, 8, 8, 8); break; // THR=2*1=2 (well under band)
108+
case 1: TN_DISPATCH(4, 8, 8, 16); break; // THR=2*2=4
109+
case 2: TN_DISPATCH(4, 8, 8, 32); break; // THR=2*4=8
110+
case 3: TN_DISPATCH(4, 8, 8, 48); break; // THR=2*6=12
111+
// BLK_M=16 rungs (nw2<=16). DIM_X=4, THR_M=4.
112+
case 4: TN_DISPATCH(4, 8, 16, 8); break; // THR=4*1=4
113+
case 5: TN_DISPATCH(4, 8, 16, 16); break; // THR=4*2=8
114+
case 6: TN_DISPATCH(4, 8, 16, 32); break; // THR=4*4=16 (in band)
115+
case 7: TN_DISPATCH(4, 8, 16, 48); break; // THR=4*6=24 (in band)
116+
// BLK_M=32 rungs (nw2<=32). DIM_X=8, THR_M=4.
117+
case 8: TN_DISPATCH(8, 4, 32, 8); break; // THR=4*2=8
118+
case 9: TN_DISPATCH(8, 4, 32, 16); break; // THR=4*4=16 (in band)
119+
case 10: TN_DISPATCH(8, 8, 32, 32); break; // THR=4*4=16 (in band)
120+
case 11: TN_DISPATCH(8, 8, 32, 48); break; // THR=4*6=24 (in band)
121+
// BLK_M=48 rungs (nw2<=48). DIM_X=8, THR_M=6.
122+
case 12: TN_DISPATCH(8, 4, 48, 8); break; // THR=6*2=12
123+
case 13: TN_DISPATCH(8, 4, 48, 16); break; // THR=6*4=24 (in band)
124+
case 14: TN_DISPATCH(8, 8, 48, 32); break; // THR=6*4=24 (in band)
125+
case 15: TN_DISPATCH(8, 8, 48, 48); break; // THR=6*6=36 (top of band)
126+
}
127+
128+
#undef TN_DISPATCH
41129
}
42130

43131
// Explicit instantiations
44132
template void gemm_nn_vbatch<double>(
45-
int, int, int, const int*, const int*, const int*,
133+
int, int, int,
46134
const double* const*, const int*, const double* const*, const int*,
47135
double**, const int*, int, cudaStream_t, const double*);
48136

49137
template void gemm_nn_vbatch<float>(
50-
int, int, int, const int*, const int*, const int*,
138+
int, int, int,
51139
const float* const*, const int*, const float* const*, const int*,
52140
double**, const int*, int, cudaStream_t, const float*);
53141

54142
template void gemm_tn_vbatch<double>(
55-
int, int, int, const int*, const int*, const int*,
143+
int, int, int,
56144
const double* const*, const int*, const double* const*, const int*,
57145
double**, const int*, int, cudaStream_t, const double*);
58146

59147
template void gemm_tn_vbatch<float>(
60-
int, int, int, const int*, const int*, const int*,
148+
int, int, int,
61149
const float* const*, const int*, const float* const*, const int*,
62150
double**, const int*, int, cudaStream_t, const float*);

‎source/source_lcao/module_gint/kernel/dgemm_vbatch.h‎

Lines changed: 19 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -2,61 +2,37 @@
22

33
#include <cuda_runtime.h>
44

5-
// Template version: C(batch_id) = alpha * A(batch_id) * B(batch_id) + C(batch_id)
6-
// As with gemm_tn_vbatch, the C accumulator is always double regardless of the
7-
// input type T so the per-block reduction and device-side atomicAdd run in fp64.
5+
// Shape-exact batched GEMM dispatchers.
6+
//
7+
// Every (A_i, B_i, C_i) in the batch has exactly the same (m, n, k); the
8+
// caller (phi_operator_gpu.cu) guarantees this by bucketing atom pairs on
9+
// (nw1, nw2) before calling in. The scalar m/n/k drive tile-ladder selection,
10+
// grid sizing, and the kernel itself -- there is no per-batch-id M/N/K
11+
// indirection.
12+
//
13+
// The C output is always double, independent of T. For T=float the per-item
14+
// inner products accumulate in fp32, but the cross-item accumulation into a
15+
// shared C element is done with a device-side fp64 atomicAdd (see the kernels'
16+
// store loop), so summing many atom-pair contributions into the same
17+
// hr_gint / phi_dm element does not drift. For T=double, A, B and C are all
18+
// double.
19+
20+
// C(batch) = alpha * A(batch) * B(batch) + C(batch)
821
template<typename T>
922
void gemm_nn_vbatch(
10-
int max_m, int max_n, int max_k,
11-
const int* m_d, const int* n_d, const int* k_d,
23+
int m, int n, int k,
1224
const T* const* A_array_d, const int* lda_d,
1325
const T* const* B_array_d, const int* ldb_d,
1426
double** C_array_d, const int* ldc_d,
1527
int batchCount, cudaStream_t stream,
1628
const T* alpha = nullptr);
1729

18-
// Template version: C(batch_id) = alpha * A(batch_id)^T * B(batch_id) + C(batch_id)
19-
// The C accumulator is always double regardless of input type T: a fp32 GEMM
20-
// path (T=float) feeds fp32 multiplies into fp64 accumulators (registers and
21-
// device-side atomicAdds) to avoid catastrophic precision loss across many
22-
// atom-pair contributions to the same hr_gint element.
30+
// C(batch) = alpha * A(batch)^T * B(batch) + C(batch)
2331
template<typename T>
2432
void gemm_tn_vbatch(
25-
int max_m, int max_n, int max_k,
26-
const int* m_d, const int* n_d, const int* k_d,
33+
int m, int n, int k,
2734
const T* const* A_array_d, const int* lda_d,
2835
const T* const* B_array_d, const int* ldb_d,
2936
double** C_array_d, const int* ldc_d,
3037
int batchCount, cudaStream_t stream,
3138
const T* alpha = nullptr);
32-
33-
// Legacy double-only aliases for backward compatibility
34-
inline void dgemm_nn_vbatch(
35-
int max_m, int max_n, int max_k,
36-
const int* m_d, const int* n_d, const int* k_d,
37-
const double* const* A_array_d, const int* lda_d,
38-
const double* const* B_array_d, const int* ldb_d,
39-
double** C_array_d, const int* ldc_d,
40-
int batchCount, cudaStream_t stream,
41-
const double* alpha = nullptr)
42-
{
43-
gemm_nn_vbatch<double>(max_m, max_n, max_k,
44-
m_d, n_d, k_d, A_array_d, lda_d, B_array_d, ldb_d,
45-
C_array_d, ldc_d, batchCount, stream, alpha);
46-
}
47-
48-
inline void dgemm_tn_vbatch(
49-
int max_m, int max_n, int max_k,
50-
const int* m_d, const int* n_d, const int* k_d,
51-
const double* const* A_array_d, const int* lda_d,
52-
const double* const* B_array_d, const int* ldb_d,
53-
double** C_array_d, const int* ldc_d,
54-
int batchCount, cudaStream_t stream,
55-
const double* alpha = nullptr)
56-
{
57-
// T=double path: A, B, and C are all double — the C-channel double-fix
58-
// matches the legacy signature here.
59-
gemm_tn_vbatch<double>(max_m, max_n, max_k,
60-
m_d, n_d, k_d, A_array_d, lda_d, B_array_d, ldb_d,
61-
C_array_d, ldc_d, batchCount, stream, alpha);
62-
}

0 commit comments

Comments
 (0)