forked from zhongkaifu/TensorSharp
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathTensorIterState.cs
More file actions
154 lines (132 loc) · 4.68 KB
/
Copy pathTensorIterState.cs
File metadata and controls
154 lines (132 loc) · 4.68 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
// Copyright (c) Zhongkai Fu. All rights reserved.
// https://github.com/zhongkaifu/TensorSharp
//
// This file is part of TensorSharp.
//
// TensorSharp is licensed under the BSD-3-Clause license found in the LICENSE file in the root directory of this source tree.
//
// TensorSharp is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the BSD-3-Clause License for more details.
using System;
namespace TensorSharp
{
public class TensorIterState
{
private readonly ReadOnlyMemory<long> sizes;
private readonly ReadOnlyMemory<long> strides;
private long stride, size;
private int dim;
private readonly long[] counter;
private readonly int step;
private long index;
unsafe public float* data;
unsafe public TensorIterState(float* buffer, int dimCount, long[] sizes, long[] strides, int step = 1)
: this(buffer, dimCount, (ReadOnlyMemory<long>)sizes, (ReadOnlyMemory<long>)strides, step)
{
}
unsafe public TensorIterState(float* buffer, int dimCount, ReadOnlyMemory<long> sizes, ReadOnlyMemory<long> strides, int step = 1)
{
if (sizes.Length < dimCount || strides.Length < dimCount)
{
throw new ArgumentException("sizes and strides must contain dimCount elements");
}
this.sizes = sizes;
this.strides = strides;
this.step = step;
ReadOnlySpan<long> sizesSpan = sizes.Span;
ReadOnlySpan<long> stridesSpan = strides.Span;
index = 0;
data = buffer;
for (dim = dimCount - 1; dim >= 0; dim--)
{
if (sizesSpan[dim] != 1)
{
break;
}
}
// Get stride for dimension
stride = (dim == -1 ? 0 : stridesSpan[dim]);
// Find largest contiguous section.
// Note: this updates dim and size.
size = 1;
for (dim = dimCount - 1; dim >= 0; dim--)
{
if (stridesSpan[dim] == size)
{
size *= sizesSpan[dim];
}
else
{
break;
}
}
if (size % step != 0)
{
throw new ArgumentException($"Size '{size}' mod step '{step}' must be zero.");
}
// Counter keeps track of dimensions outside the contiguous block.
counter = new long[dim + 1];
for (int i = 0; i < dim + 1; ++i)
{
counter[i] = 0;
}
}
public bool ReachedBlockEnd()
{
return !(index < size);
}
public void BlockStep()
{
unsafe
{
index += step;
data += (stride * step);
}
}
// Returns true if there is another block to iterate over,
// returns false if we are at end of iteration
public bool NextBlock()
{
unsafe
{
// If not at end of current block yet, do nothing
if (index == size)
{
// If contiguous block encompassed all dimensions, we are done
if (dim == -1)
{
return false;
}
ReadOnlySpan<long> sizesSpan = sizes.Span;
ReadOnlySpan<long> stridesSpan = strides.Span;
// Reset data offset
data -= size * stride;
// Update counter and data for next contiguous block
for (long j = dim; j >= 0; --j)
{
counter[j]++;
data += stridesSpan[(int)j];
if (counter[j] == sizesSpan[(int)j])
{
if (j == 0)
{
return false;
}
else
{
data -= counter[j] * stridesSpan[(int)j];
counter[j] = 0;
}
}
else
{
break;
}
}
index = 0;
}
return true;
}
}
}
}