Skip to content

Commit f385af7

Browse files
Merge pull request #54 from PanZezhong1725/issue/53
issue/53: 形状统一为size_t,步长统一为ptrdiff_t
2 parents 3c31dc6 + 64849b4 commit f385af7

File tree

19 files changed

+59
-59
lines changed

19 files changed

+59
-59
lines changed

include/infiniop/ops/attention.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ __C __export infiniopStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t
1515
infiniopTensorDescriptor_t v_desc,
1616
infiniopTensorDescriptor_t k_cache_desc,
1717
infiniopTensorDescriptor_t v_cache_desc,
18-
uint64_t pos);
18+
size_t pos);
1919

2020
__C __export infiniopStatus_t infiniopGetAttentionWorkspaceSize(infiniopAttentionDescriptor_t desc, size_t *size);
2121

include/infiniop/ops/avg_pool.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,10 @@ __C __export infiniopStatus_t infiniopCreateAvgPoolDescriptor(infiniopHandle_t h
99
infiniopAvgPoolDescriptor_t *desc_ptr,
1010
infiniopTensorDescriptor_t y,
1111
infiniopTensorDescriptor_t x,
12-
uint64_t const *kernel_shape,
13-
uint64_t const *pads,
14-
int64_t const *strides,
15-
uint64_t n);
12+
size_t const *kernel_shape,
13+
size_t const *pads,
14+
ptrdiff_t const *strides,
15+
size_t n);
1616

1717
__C __export infiniopStatus_t infiniopGetAvgPoolWorkspaceSize(infiniopAvgPoolDescriptor_t desc, size_t *size);
1818

include/infiniop/ops/max_pool.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,10 @@ __C __export infiniopStatus_t infiniopCreateMaxPoolDescriptor(infiniopHandle_t h
99
infiniopMaxPoolDescriptor_t *desc_ptr,
1010
infiniopTensorDescriptor_t y,
1111
infiniopTensorDescriptor_t x,
12-
uint64_t const *kernel_shape,
13-
uint64_t const *pads,
14-
int64_t const *strides,
15-
uint64_t n);
12+
size_t const *kernel_shape,
13+
size_t const *pads,
14+
ptrdiff_t const *strides,
15+
size_t n);
1616

1717
__C __export infiniopStatus_t infiniopGetMaxPoolWorkspaceSize(infiniopMaxPoolDescriptor_t desc, size_t *size);
1818

include/infiniop/tensor_descriptor.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ struct InfiniopTensorDescriptor {
1212
// Shape of the tensor, ndim elements
1313
size_t *shape;
1414
// Stride of each dimension in elements, ndim elements
15-
int64_t *strides;
15+
ptrdiff_t *strides;
1616
};
1717

1818
typedef struct InfiniopTensorDescriptor *infiniopTensorDescriptor_t;
1919

20-
__C __export infiniopStatus_t infiniopCreateTensorDescriptor(infiniopTensorDescriptor_t *desc_ptr, size_t ndim, size_t const *shape, int64_t const *strides, infiniDtype_t dtype);
20+
__C __export infiniopStatus_t infiniopCreateTensorDescriptor(infiniopTensorDescriptor_t *desc_ptr, size_t ndim, size_t const *shape, ptrdiff_t const *strides, infiniDtype_t dtype);
2121

2222
__C __export infiniopStatus_t infiniopDestroyTensorDescriptor(infiniopTensorDescriptor_t desc);
2323

src/infiniop/devices/cpu/common_cpu.cc

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,8 @@ uint16_t f32_to_f16(float val) {
6060
}
6161

6262
size_t indexToReducedOffset(size_t flat_index, size_t ndim,
63-
int64_t const *broadcasted_strides,
64-
int64_t const *target_strides) {
63+
ptrdiff_t const *broadcasted_strides,
64+
ptrdiff_t const *target_strides) {
6565
size_t res = 0;
6666
for (size_t i = 0; i < ndim; ++i) {
6767
res += flat_index / broadcasted_strides[i] * target_strides[i];
@@ -71,7 +71,7 @@ size_t indexToReducedOffset(size_t flat_index, size_t ndim,
7171
}
7272

7373
size_t indexToOffset(size_t flat_index, size_t ndim, size_t const *shape,
74-
int64_t const *strides) {
74+
ptrdiff_t const *strides) {
7575
size_t res = 0;
7676
for (size_t i = ndim; i-- >= 0;) {
7777
res += (flat_index % shape[i]) * strides[i];
@@ -81,7 +81,7 @@ size_t indexToOffset(size_t flat_index, size_t ndim, size_t const *shape,
8181
}
8282

8383
size_t getPaddedSize(size_t ndim, size_t *shape, size_t const *pads) {
84-
uint64_t total_size = 1;
84+
size_t total_size = 1;
8585
for (size_t i = 0; i < ndim; ++i) {
8686
total_size *= shape[i] + (i < 2 ? 0 : 2 * pads[i - 2]);
8787
}

src/infiniop/devices/cpu/common_cpu.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,10 @@ float f16_to_f32(uint16_t code);
1313
uint16_t f32_to_f16(float val);
1414

1515
// return the memory offset of original tensor, given the flattened index of broadcasted tensor
16-
size_t indexToReducedOffset(size_t flat_index, size_t ndim, int64_t const *broadcasted_strides, int64_t const *target_strides);
16+
size_t indexToReducedOffset(size_t flat_index, size_t ndim, ptrdiff_t const *broadcasted_strides, ptrdiff_t const *target_strides);
1717

1818
// return the memory offset a tensor given flattened index
19-
size_t indexToOffset(size_t flat_index, size_t ndim, size_t const *shape, int64_t const *strides);
19+
size_t indexToOffset(size_t flat_index, size_t ndim, size_t const *shape, ptrdiff_t const *strides);
2020

2121
/**
2222
* get the total array size (element count) after applying padding for a

src/infiniop/devices/cuda/common_cuda.cuh

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,8 @@ inline cudnnDataType_t getCudnnDtype(infiniDtype_t dt) {
9696
// return the memory offset of original tensor, given the flattened index of
9797
// broadcasted tensor
9898
inline __device__ __host__ size_t indexToReducedOffset(
99-
size_t flat_index, size_t ndim, int64_t const *broadcasted_strides,
100-
int64_t const *target_strides) {
99+
size_t flat_index, size_t ndim, ptrdiff_t const *broadcasted_strides,
100+
ptrdiff_t const *target_strides) {
101101
size_t res = 0;
102102
for (size_t i = 0; i < ndim; ++i) {
103103
res += flat_index / broadcasted_strides[i] * target_strides[i];
@@ -109,7 +109,7 @@ inline __device__ __host__ size_t indexToReducedOffset(
109109
// get the memory offset of the given element in a tensor given its flat index
110110
inline __device__ __host__ size_t indexToOffset(size_t flat_index, size_t ndim,
111111
size_t const *shape,
112-
int64_t const *strides) {
112+
ptrdiff_t const *strides) {
113113
size_t res = 0;
114114
for (size_t i = ndim; i-- > 0;) {
115115
res += (flat_index % shape[i]) * strides[i];

src/infiniop/ops/causal_softmax/operator.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ __C infiniopStatus_t infiniopCreateCausalSoftmaxDescriptor(
4040
return INFINIOP_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
4141
}
4242

43-
__C infiniopStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDescriptor_t desc, uint64_t *size) {
43+
__C infiniopStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDescriptor_t desc, size_t *size) {
4444
switch (desc->device) {
4545
#ifdef ENABLE_CPU
4646
case DevCpu:
@@ -78,7 +78,7 @@ __C infiniopStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmax
7878
return INFINIOP_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
7979
}
8080

81-
__C infiniopStatus_t infiniopCausalSoftmax(infiniopCausalSoftmaxDescriptor_t desc, void *workspace, uint64_t workspace_size, void *data, void *stream) {
81+
__C infiniopStatus_t infiniopCausalSoftmax(infiniopCausalSoftmaxDescriptor_t desc, void *workspace, size_t workspace_size, void *data, void *stream) {
8282
switch (desc->device) {
8383
#ifdef ENABLE_CPU
8484
case DevCpu:

src/infiniop/ops/matmul/blas.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,11 @@
99
typedef struct BlasMatrix {
1010
size_t ndim;
1111
size_t batch;
12-
int64_t stride;
12+
ptrdiff_t stride;
1313
size_t rows;
1414
size_t cols;
15-
int64_t row_stride;
16-
int64_t col_stride;
15+
ptrdiff_t row_stride;
16+
ptrdiff_t col_stride;
1717

1818
BlasMatrix() {}
1919

@@ -56,7 +56,7 @@ typedef struct BlasMatrix {
5656
std::swap(row_stride, col_stride);
5757
}
5858

59-
int64_t ld() const {
59+
ptrdiff_t ld() const {
6060
if (this->row_stride == 1) {
6161
return this->col_stride;
6262
} else {

src/infiniop/ops/matmul/cpu/matmul_cpu.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ infiniopStatus_t cpuCreateMatmulDescriptor(
2525
}
2626

2727
infiniopStatus_t cpuGetMatmulWorkspaceSize(infiniopMatmulCpuDescriptor_t desc,
28-
uint64_t *size) {
28+
size_t *size) {
2929
*size = 0;
3030
return INFINIOP_STATUS_SUCCESS;
3131
}
@@ -76,7 +76,7 @@ infiniopStatus_t cpuCalculateMatmul(infiniopMatmulCpuDescriptor_t desc, void *c,
7676
}
7777

7878
infiniopStatus_t cpuMatmul(infiniopMatmulCpuDescriptor_t desc, void *workspace,
79-
uint64_t workspace_size, void *c, void const *a,
79+
size_t workspace_size, void *c, void const *a,
8080
void const *b, float alpha, float beta) {
8181
if (desc->dtype == INFINI_DTYPE_F16) {
8282
return cpuCalculateMatmul<uint16_t>(desc, c, beta, a, b, alpha);

0 commit comments

Comments
 (0)