Repository navigation
Expand file tree
/
Copy pathNDIter.cs
More file actions
4782 lines (4247 loc) · 219 KB
/
Copy pathNDIter.cs
File metadata and controls
4782 lines (4247 loc) · 219 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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
using System;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using NumSharp.Backends.Kernels;
using NumSharp.Utilities;
namespace NumSharp.Backends.Iteration
{
/// <summary>
/// Function to advance iterator to next position.
/// Returns true if more iterations remain.
/// </summary>
public unsafe delegate bool NDIterNextFunc(ref NDIterState state);
/// <summary>
/// Function to get multi-index at current position.
/// </summary>
public unsafe delegate void NDIterGetMultiIndexFunc(ref NDIterState state, long* outCoords);
/// <summary>
/// Inner loop kernel called by iterator.
/// </summary>
public unsafe delegate void NDIterInnerLoopFunc(
void** dataptrs,
long* strides,
long count,
void* auxdata);
/// <summary>
/// High-performance multi-operand iterator matching NumPy's nditer API.
/// </summary>
public unsafe ref partial struct NDIterRef
{
private NDIterState* _state;
private bool _ownsState;
// Borrowed: an iterator walks the caller's operands and never owns them (NumPy's nditer
// contract). The one array it substitutes itself — a COPY_IF_OVERLAP temporary — is written
// back and released by ResolveWritebacks, not by any owner-of-the-slot rule. (An ALLOCATE
// output is yielded to the caller through the operands array.)
[NDBorrowed] private NDArray[]? _operands;
private NDIterNextFunc? _cachedIterNext;
/// <summary>
/// Write-back registrations created by COPY_IF_OVERLAP: when entry i is
/// non-null, <see cref="_operands"/>[i] is a temporary copy and the
/// stored array is the user's original operand. <see cref="Dispose"/>
/// copies the temporary back (NumPy's WRITEBACKIFCOPY resolved at
/// NDIter_Deallocate). Not duplicated by <see cref="Copy"/> and not
/// carried by <c>TransferStateOwnership</c> — the write-back belongs to
/// the constructing NDIterRef. Borrowed: the entries ARE the user's operands.
/// </summary>
[NDBorrowed] private NDArray?[]? _writebackOriginals;
// =========================================================================
// Factory Methods
// =========================================================================
/// <summary>
/// Create single-operand iterator.
/// Equivalent to NumPy's NDIter_New.
/// </summary>
public static NDIterRef New(
NDArray op,
NDIterGlobalFlags flags = NDIterGlobalFlags.None,
NPY_ORDER order = NPY_ORDER.NPY_KEEPORDER,
NPY_CASTING casting = NPY_CASTING.NPY_SAFE_CASTING,
NPTypeCode? dtype = null)
{
var opFlags = new[] { NDIterPerOpFlags.READONLY };
var dtypes = dtype.HasValue ? new[] { dtype.Value } : null;
return MultiNew(1, new[] { op }, flags, order, casting, opFlags, dtypes);
}
/// <summary>
/// Create multi-operand iterator.
/// Equivalent to NumPy's NDIter_MultiNew.
/// </summary>
public static NDIterRef MultiNew(
int nop,
NDArray[] op,
NDIterGlobalFlags flags,
NPY_ORDER order,
NPY_CASTING casting,
NDIterPerOpFlags[] opFlags,
NPTypeCode[]? opDtypes = null)
{
return AdvancedNew(nop, op, flags, order, casting, opFlags, opDtypes);
}
/// <summary>
/// Create iterator with full control over all parameters.
/// Equivalent to NumPy's NDIter_AdvancedNew.
/// </summary>
public static NDIterRef AdvancedNew(
int nop,
NDArray[] op,
NDIterGlobalFlags flags,
NPY_ORDER order,
NPY_CASTING casting,
NDIterPerOpFlags[] opFlags,
NPTypeCode[]? opDtypes = null,
int opAxesNDim = -1,
int[][]? opAxes = null,
long[]? iterShape = null,
long bufferSize = 0)
{
if (nop < 1)
throw new ArgumentOutOfRangeException(nameof(nop), "At least one operand is required");
if (op == null || op.Length < nop)
throw new ArgumentException("Operand array must contain at least nop elements", nameof(op));
if (opFlags == null || opFlags.Length < nop)
throw new ArgumentException("OpFlags array must contain at least nop elements", nameof(opFlags));
// Allocate state on heap for ref struct lifetime — one block carrying the header and
// the arena the dimension/operand arrays are carved from (recycled per thread).
var statePtr = AllocateStateBlock();
try
{
var iter = new NDIterRef
{
_state = statePtr,
_ownsState = true,
_operands = op,
};
iter.Initialize(nop, op, flags, order, casting, opFlags, opDtypes, opAxesNDim, opAxes, iterShape, bufferSize);
return iter;
}
catch
{
// Tear down whatever Initialize managed to build (buffers, external dimension
// arrays) and hand the block back to the cache / allocator.
FreeState(statePtr);
throw;
}
}
// =========================================================================
// State block allocation — single block + per-thread recycling
// =========================================================================
//
// Every heap NDIterState is ONE native allocation: the header followed by
// StateArenaBytes of arena that NDIterState.AllocateDimArrays carves its dimension
// and per-operand arrays from (see the arena notes in NDIter.State.cs). Freed blocks
// are parked in a small per-thread cache and handed to the next construction after
// ResetForRecycle re-zeroed only what the previous iterator used. Measured on the
// three-operand 1-D construction: 161 ns → ~90 ns; the three calloc/free pairs it
// replaces were 78 ns of that (one 616-byte calloc/free is 29 ns). The cache is
// bounded (StateBlockCacheSlots blocks per thread, ~1.9 KB each) and blocks are
// uniform, so there is no size-class logic; a block freed on another thread simply
// joins that thread's cache. States allocated elsewhere (no arena) are freed
// directly, so ReleaseStateBlock is safe for any pointer FreeState accepts.
/// <summary>Bytes reserved after the state header for the dimension + per-operand arrays.</summary>
internal const int StateArenaBytes = 1536;
/// <summary>Blocks parked per thread by <see cref="ReleaseStateBlock"/>.</summary>
internal const int StateBlockCacheSlots = 4;
private sealed class StateBlockCache
{
public readonly nint[] Slots = new nint[StateBlockCacheSlots];
public int Count;
}
[ThreadStatic]
private static StateBlockCache? t_stateBlocks;
/// <summary>
/// Obtain a zeroed state header with an attached, zeroed inline arena — from the
/// per-thread cache when one is parked there, else a fresh single allocation.
/// </summary>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static NDIterState* AllocateStateBlock()
{
var cache = t_stateBlocks;
if (cache is not null && cache.Count > 0)
return (NDIterState*)cache.Slots[--cache.Count];
var block = (NDIterState*)NativeMemory.AllocZeroed((nuint)(sizeof(NDIterState) + StateArenaBytes));
block->AttachInlineArena((byte*)block + sizeof(NDIterState), StateArenaBytes);
return block;
}
/// <summary>
/// Release a state header whose buffers and external dimension arrays have already
/// been freed: single-block states are re-zeroed and parked in the per-thread cache
/// (or freed when it is full); states without an inline arena are freed directly.
/// </summary>
internal static void ReleaseStateBlock(NDIterState* state)
{
if (state == null)
return;
if (!state->HasInlineArena)
{
NativeMemory.Free(state);
return;
}
var cache = t_stateBlocks ??= new StateBlockCache();
if (cache.Count < StateBlockCacheSlots)
{
state->ResetForRecycle();
cache.Slots[cache.Count++] = (nint)state;
return;
}
NativeMemory.Free(state);
}
/// <summary>Blocks currently parked in the calling thread's cache (diagnostics/tests).</summary>
internal static int CachedStateBlockCount => t_stateBlocks?.Count ?? 0;
[MethodImpl(MethodImplOptions.AggressiveOptimization)]
private void Initialize(
int nop,
NDArray[] op,
NDIterGlobalFlags flags,
NPY_ORDER order,
NPY_CASTING casting,
NDIterPerOpFlags[] opFlags,
NPTypeCode[]? opDtypes,
int opAxesNDim,
int[][]? opAxes,
long[]? iterShape,
long bufferSize)
{
_state->MaskOp = -1;
_state->IterStart = 0;
// Pre-check per-operand flags BEFORE allocation (nop arg, not state):
// WRITEMASKED/ARRAYMASK pairing, VIRTUAL/null-operand rules.
// The actual MaskOp assignment happens after AllocateDimArrays when NOp is set.
if (opFlags != null)
PreCheckMaskOpPairing(nop, op, opFlags);
// Calculate broadcast shape, optionally overridden by iterShape.
// Dimensions are long (NumPy npy_intp parity) so large axes (> int.MaxValue
// elements) survive iterator construction without narrowing.
long[] broadcastShape;
if (iterShape != null && iterShape.Length > 0)
{
// Use explicit iterShape - allows specifying iteration shape different from broadcast
// NumPy's NDIter_AdvancedNew() uses this for reductions and custom iteration patterns
broadcastShape = (long[])iterShape.Clone();
// Validate that operands are compatible with the specified shape
// Pass opAxes so validation accounts for -1 entries (broadcast/reduce axes)
ValidateIterShape(nop, op, opFlags, broadcastShape, opAxesNDim, opAxes);
}
else
{
broadcastShape = CalculateBroadcastShape(nop, op, opFlags, opAxesNDim, opAxes);
// Validate NO_BROADCAST operands match without stretching
ValidateIterShape(nop, op, opFlags, broadcastShape, opAxesNDim, opAxes);
}
// COPY_IF_OVERLAP: if any write operand may share memory with any
// read operand, replace the WRITE operand with a temporary copy and
// register a write-back to the original on Dispose. This is what
// makes ufunc-style calls with overlapping in/out well-defined
// (e.g. add(a[:-1], a[:-1], out=a[1:])). Runs before the ALLOCATE
// loop so iterator-allocated outputs (null here) are skipped — a
// fresh allocation can never overlap an input.
// NumPy: nditer_constr.c:3083-3168 (npyiter_allocate_arrays).
if ((flags & NDIterGlobalFlags.COPY_IF_OVERLAP) != 0)
{
var forceCopy = ComputeCopyIfOverlap(nop, op, opFlags);
if (forceCopy != null)
MaterializeForcedCopies(nop, op, opFlags, forceCopy);
}
// Allocate null operands that have ALLOCATE or VIRTUAL set.
// NumPy: npyiter_allocate_arrays in nditer_constr.c allocates a temp
// array for EVERY null operand — including VIRTUAL ones, whose
// NPY_OP_ITFLAG_VIRTUAL is consumed nowhere but DebugPrint (the
// NEP-12 buffer-only semantics never landed). Verified NumPy 2.4.2:
// a virtual operand gets a real backing array; writes through the
// iterator land in it; it.operands[i] exposes it.
// Allocated output has shape = broadcastShape (accounting for op_axes).
// Dtype: ALLOCATE takes opDtypes[opIdx] (NumSharp requires it);
// VIRTUAL-only DISCARDS any requested dtype (npyiter_prepare_one_operand
// c:1037-1039 nulls op_dtype for the non-ALLOCATE branch) and resolves
// the common dtype of the real operands instead — except an ARRAYMASK
// operand, which defaults to bool (c:1041-1049). When both flags are
// set, ALLOCATE wins (c:1009 checks it first) and the request holds.
for (int opIdx = 0; opIdx < nop; opIdx++)
{
if (op[opIdx] is null &&
(opFlags[opIdx] & (NDIterPerOpFlags.ALLOCATE | NDIterPerOpFlags.VIRTUAL)) != 0)
{
NPTypeCode allocDtype;
bool isArrayMaskOp = (opFlags[opIdx] & NDIterPerOpFlags.ARRAYMASK) != 0;
if ((opFlags[opIdx] & NDIterPerOpFlags.ALLOCATE) != 0)
{
bool hasRequest = opDtypes != null && opIdx < opDtypes.Length &&
opDtypes[opIdx] != NPTypeCode.Empty;
if (hasRequest)
allocDtype = opDtypes[opIdx];
else if (isArrayMaskOp)
allocDtype = NPTypeCode.Boolean; // mask default (c:1041-1049)
else
throw new ArgumentException(
$"Operand {opIdx} is null with ALLOCATE flag but opDtypes is not provided", nameof(opDtypes));
}
else if (isArrayMaskOp)
{
allocDtype = NPTypeCode.Boolean; // virtual mask default (c:1041-1049)
}
else
{
// VIRTUAL: common dtype over the real operands (the
// requested dtype was discarded, NumPy parity).
allocDtype = NPTypeCode.Empty;
for (int j = 0; j < nop; j++)
{
if (op[j] is null)
continue;
allocDtype = allocDtype == NPTypeCode.Empty
? op[j].typecode
: NDIterCasting.PromoteTypes(allocDtype, op[j].typecode);
}
if (allocDtype == NPTypeCode.Empty)
throw new ArgumentException(
"at least one array or dtype is required");
}
// Determine output shape: for op_axes, filter out -1 entries
long[] outputShape;
if (opAxes != null && opIdx < opAxes.Length && opAxes[opIdx] != null)
{
var axisMap = opAxes[opIdx];
// Count non-negative entries
int realNDim = 0;
for (int i = 0; i < axisMap.Length; i++)
if (axisMap[i] >= 0) realNDim++;
outputShape = new long[realNDim];
int outIdx = 0;
for (int iterAxis = 0; iterAxis < axisMap.Length && iterAxis < broadcastShape.Length; iterAxis++)
{
// Non-negative entries map iterShape axes to output axes
if (axisMap[iterAxis] >= 0)
outputShape[axisMap[iterAxis]] = broadcastShape[iterAxis];
// -1 entries are "reduced" dimensions - not in output shape
}
}
else
{
// No op_axes: output has full broadcast shape
outputShape = (long[])broadcastShape.Clone();
}
// Allocate the NDArray with the resolved dtype and shape
var shape = outputShape.Length == 0 ? new Shape() : new Shape(outputShape);
op[opIdx] = np.zeros(shape, allocDtype);
}
}
// Update _operands so it reflects the allocated arrays
_operands = op;
// =========================================================================
// NUMSHARP DIVERGENCE: Allocate dimension arrays dynamically
// Unlike NumPy's fixed NPY_MAXDIMS=64, NumSharp supports unlimited dimensions.
// Arrays are allocated based on actual ndim for memory efficiency.
// =========================================================================
_state->AllocateDimArrays(broadcastShape.Length, nop);
// Set IDENTPERM on construction. Perm starts as identity (set by AllocateDimArrays);
// reordering (ReorderAxesForCoalescing) and flipping (FlipNegativeStrides) clear
// this flag when they mutate perm. Matches NumPy nditer_constr.c:262-264.
_state->ItFlags |= (uint)NDIterFlags.IDENTPERM;
// Set MaskOp for ARRAYMASK operand (if any). Requires NOp to be set by
// AllocateDimArrays above. NumPy nditer_constr.c:1184-1196.
if (opFlags != null)
SetMaskOpFromFlags(opFlags);
_state->IterSize = 1;
for (int d = 0; d < _state->NDim; d++)
{
_state->Shape[d] = broadcastShape[d];
_state->IterSize *= broadcastShape[d];
}
_state->IterEnd = _state->IterSize;
// Handle zero-size iteration
if (_state->IterSize == 0 && (flags & NDIterGlobalFlags.ZEROSIZE_OK) == 0)
{
// Just allow it anyway for now
}
// Determine common dtype if COMMON_DTYPE flag is set
NPTypeCode? commonDtype = null;
if ((flags & NDIterGlobalFlags.COMMON_DTYPE) != 0)
{
commonDtype = NDIterCasting.FindCommonDtype(op, nop);
}
// Set up operands
bool anyNeedsCast = false;
for (int i = 0; i < nop; i++)
{
var arr = op[i];
var arrShape = arr.Shape;
// Store source dtype (actual array dtype)
_state->SetOpSrcDType(i, arr.typecode);
// Determine buffer/target dtype. virtual-without-allocate is a
// single masked equality so the common path pays one AND.
NPTypeCode bufferDtype;
bool virtualAllocated =
(opFlags[i] & (NDIterPerOpFlags.VIRTUAL | NDIterPerOpFlags.ALLOCATE))
== NDIterPerOpFlags.VIRTUAL;
if (virtualAllocated)
{
// VIRTUAL discarded any requested dtype at allocation
// (NumPy parity — see the allocation loop above): the
// freshly allocated array IS the iteration dtype, no cast.
bufferDtype = arr.typecode;
}
else if (opDtypes != null && i < opDtypes.Length && opDtypes[i] != NPTypeCode.Empty)
{
bufferDtype = opDtypes[i];
}
else if (commonDtype.HasValue)
{
bufferDtype = commonDtype.Value;
}
else
{
bufferDtype = arr.typecode;
}
_state->SetOpDType(i, bufferDtype);
// Track if any operand needs casting
if (arr.typecode != bufferDtype)
{
anyNeedsCast = true;
}
// Set operand flags
var opFlag = TranslateOpFlags(opFlags[i]);
// If operand needs casting, add CAST flag
if (arr.typecode != bufferDtype)
{
opFlag |= NDIterOpFlags.CAST;
}
// COPY_IF_OVERLAP replaced this operand with a temporary that
// is written back on Dispose (NumPy: FORCECOPY + WRITEBACKIFCOPY).
if (_writebackOriginals?[i] is not null)
{
opFlag |= NDIterOpFlags.FORCECOPY | NDIterOpFlags.HAS_WRITEBACK;
}
_state->SetOpFlags(i, opFlag);
// Calculate strides for this operand
var stridePtr = _state->GetStridesPointer(i);
byte* basePtr;
// Check if op_axes is provided for this operand
if (opAxes != null && i < opAxes.Length && opAxes[i] != null)
{
// Use op_axes mapping to set up strides directly.
// NumPy fill_axisdata parity (nditer_constr.c:1618-1679):
// the stride is 0 when the iterator dim is 1, when the entry
// is a broadcast/reduction axis (-1 or REDUCTION_AXIS-encoded),
// or when the operand's mapped axis has length 1 (broadcast
// stretch) — using the raw array stride for a stretched size-1
// axis walked out of bounds.
var opAxisMap = opAxes[i];
var arrStrides = arrShape.strides;
var arrDims = arrShape.dimensions;
// Operand base must point to the view's LOGICAL origin: add
// Shape.offset (in elements). The standard-broadcast branch below
// already does this (lines ~473/481); the op_axes branch did not,
// so any op_axes reduction over a view whose offset lives in
// Shape.offset (sliced a[1:3,1:3], negative-stride a[::-1] / a[:,::-1])
// read from the buffer base instead — wrong cells, and after
// FlipNegativeStrides moved the pointer, out-of-bounds reads/writes
// (garbage / NaN). Views with offset==0 (C-/F-contig, transpose,
// positive strided) were unaffected, which is why it hid so long.
// Use NPTypeCode.SizeOf() (1 byte for bool) NOT arr.dtypesize.
basePtr = (byte*)arr.Address + (arrShape.offset * arr.GetTypeCode.SizeOf());
for (int d = 0; d < _state->NDim; d++)
{
if (_state->Shape[d] == 1)
{
// Iterator dim of size 1: never advanced; NumPy forces
// stride 0 (the coalesce trivial-branch and merge rule
// rely on this invariant).
stridePtr[d] = 0;
}
else if (d < opAxisMap.Length)
{
int opAxis = NDIterUtils.GetOpAxis(opAxisMap[d], out _);
if (opAxis < 0)
{
// -1 means broadcast/reduce this dimension
stridePtr[d] = 0;
}
else if (opAxis < arrStrides.Length)
{
// Operand axis of length 1 is broadcast-stretched
// to the iter dim: stride 0, not the raw stride.
stridePtr[d] = arrDims[opAxis] == 1 ? 0 : arrStrides[opAxis];
}
else
{
stridePtr[d] = 0;
}
}
else
{
stridePtr[d] = 0;
}
}
}
else
{
// Standard broadcasting.
//
// NOTE: must use NPTypeCode.SizeOf() (1 byte for bool) and
// NOT arr.dtypesize, which is implemented via
// Marshal.SizeOf and returns 4 for bool because bool is
// marshaled as win32 BOOL. In-memory layout uses 1 byte
// per bool element, so Marshal-based sizing produces
// pointer offsets 4x too large.
int elemBytes = arr.GetTypeCode.SizeOf();
// Fast path (per-call construction cost): when the operand's
// shape already equals the iteration shape, np.broadcast_to is
// an identity transform — same offset, same strides, no stride-0
// insertion. Detect that and copy the operand's own offset/strides
// directly, skipping the broadcast_to call and its Shape allocation.
// This is the dominant elementwise case (a OP b, both same shape)
// and was ~25% of the iterator's setup overhead.
bool sameAsIter = arrShape.NDim == _state->NDim;
if (sameAsIter)
{
var adims = arrShape.dimensions;
for (int d = 0; d < _state->NDim; d++)
if (adims[d] != _state->Shape[d]) { sameAsIter = false; break; }
}
// NumPy fill_axisdata parity (nditer_constr.c:1594-1615): any
// iterator axis of size 1 gets stride 0 for every operand. The
// raw stride of a size-1 axis is arbitrary (transposed/sliced
// views carry nonzero values there) and must never survive
// into the iterator — CoalesceAxes' trivial-branch and merge
// rule both rely on the size-1 ⇒ stride-0 invariant.
if (sameAsIter)
{
basePtr = (byte*)arr.Address + (arrShape.offset * elemBytes);
var astrides = arrShape.strides;
for (int d = 0; d < _state->NDim; d++)
stridePtr[d] = _state->Shape[d] == 1 ? 0 : astrides[d];
}
else
{
var broadcastArr = np.broadcast_to(arrShape, new Shape(broadcastShape));
basePtr = (byte*)arr.Address + (broadcastArr.offset * elemBytes);
for (int d = 0; d < _state->NDim; d++)
{
stridePtr[d] = _state->Shape[d] == 1 ? 0 : broadcastArr.strides[d];
}
}
}
_state->SetDataPtr(i, basePtr);
_state->SetResetDataPtr(i, basePtr);
// Scan stretched dimensions (stride 0 with iter dim > 1).
// Read operands: record SourceBroadcast (legacy layout flag).
// Write operands: a stretched dim means several iterations land
// on the same element — that is a reduction. NumPy fill_axisdata
// (nditer_constr.c:1601-1611, npyiter_check_reduce_ok_and_set_flags)
// refuses it unless REDUCE_OK is set and the operand is readable;
// NumSharp previously filled stride 0 silently and let the
// colliding writes race. op_axes-mapped operands get the same
// treatment in ApplyOpAxes, which owns reduction semantics for
// mapped axes (including REDUCTION_AXIS-encoded entries).
bool opAxesMapped = opAxes != null && i < opAxes.Length && opAxes[i] != null;
for (int d = 0; d < _state->NDim; d++)
{
if (_state->Shape[d] > 1 && stridePtr[d] == 0)
{
if (!opAxesMapped && (opFlag & NDIterOpFlags.WRITE) != 0)
{
if ((flags & NDIterGlobalFlags.REDUCE_OK) == 0)
{
throw new ArgumentException(
$"output operand requires a reduction along dimension {d}, " +
"but the reduction is not enabled. The dimension size of 1 " +
"does not match the expected output shape.");
}
if ((opFlag & NDIterOpFlags.READ) == 0)
{
throw new ArgumentException(
"output operand requires a reduction, but is flagged as " +
"write-only, not read-write");
}
// The ARRAYMASK operand may never be the result of a
// reduction: a slot could be written while the mask
// reads True, then the reduction overwrite the mask
// to False — retroactively invalidating the write and
// violating the strict masking semantics. NumPy
// npyiter_check_reduce_ok_and_set_flags
// (nditer_constr.c:1404-1421).
if (i == _state->MaskOp)
{
throw new ArgumentException(
"output operand requires a " +
"reduction, but is flagged as " +
"the ARRAYMASK operand which " +
"is not permitted to be the " +
"result of a reduction");
}
opFlag |= NDIterOpFlags.REDUCE;
_state->SetOpFlags(i, opFlag);
_state->ItFlags |= (uint)NDIterFlags.REDUCE;
}
else
{
_state->ItFlags |= (uint)NDIterFlags.SourceBroadcast;
}
}
}
}
// Validate that casting requires BUFFERED flag
if (anyNeedsCast && (flags & NDIterGlobalFlags.BUFFERED) == 0)
{
throw new ArgumentException(
"Casting between different dtypes requires the BUFFERED flag. " +
"Add NDIterGlobalFlags.BUFFERED to enable type conversion.");
}
// Validate casting rules
NDIterCasting.ValidateCasts(ref *_state, casting);
// Apply op_axes remapping if provided
if (opAxes != null && opAxesNDim >= 0)
{
ApplyOpAxes(opAxesNDim, opAxes, flags);
}
// Deferred WRITEMASKED + REDUCE validation: every operand's strides
// are final here (standard fill AND op_axes both done), so the
// per-dim "one mask value per WRITEMASKED element" invariant can be
// checked against the mask's strides regardless of operand order.
// NumPy runs the identical loop at the tail of
// npyiter_allocate_arrays (nditer_constr.c:3351-3370).
if (_state->MaskOp >= 0)
{
for (int iop = 0; iop < nop; iop++)
{
if ((_state->GetOpFlags(iop) &
(NDIterOpFlags.WRITEMASKED | NDIterOpFlags.REDUCE)) ==
(NDIterOpFlags.WRITEMASKED | NDIterOpFlags.REDUCE))
{
CheckMaskForWriteMaskedReduction(iop);
}
}
}
// Apply axis reordering based on iteration order.
// NumPy reorders axes based on the order parameter, then coalesces if MULTI_INDEX is not set.
//
// Order semantics:
// - C-order: last axis innermost (row-major logical iteration)
// - F-order: first axis innermost (column-major logical iteration)
// - K-order: smallest stride innermost (memory-order iteration)
//
// When MULTI_INDEX is set:
// - Axes are reordered for the specified iteration order
// - No coalescing (would invalidate multi-index tracking)
// - GetMultiIndex/GotoMultiIndex use Perm to map between internal and original coords
//
// When MULTI_INDEX is NOT set:
// - Axes are reordered AND coalesced for maximum efficiency
bool hasMultiIndex = (flags & NDIterGlobalFlags.MULTI_INDEX) != 0;
// HASINDEX (C_INDEX or F_INDEX): need original axis structure preserved
// to compute the flat index correctly. Coalescing loses this info.
bool hasFlatIndex = (flags & (NDIterGlobalFlags.C_INDEX | NDIterGlobalFlags.F_INDEX)) != 0;
// Step 0: Flip negative strides for memory-order iteration
// NumPy's npyiter_flip_negative_strides() (nditer_constr.c:297-307):
// if (!(itflags & NPY_ITFLAG_FORCEDORDER)) {
// if (!any_allocate && !(flags & NPY_ITER_DONT_NEGATE_STRIDES)) {
// npyiter_flip_negative_strides(iter);
// }
// }
//
// Only K-order does NOT set FORCEDORDER. C, F, and A orders all set FORCEDORDER
// (see npyiter_apply_forced_iteration_order in nditer_constr.c:2490).
// So negative strides should only be flipped for K-order.
//
// User-visible behavior:
// - K-order on reversed array: iterate in memory order (faster)
// - C/F/A order on reversed array: iterate in logical order (user asked for it)
bool isForcedOrder = order == NPY_ORDER.NPY_CORDER
|| order == NPY_ORDER.NPY_FORTRANORDER
|| order == NPY_ORDER.NPY_ANYORDER;
if (!isForcedOrder && (flags & NDIterGlobalFlags.DONT_NEGATE_STRIDES) == 0)
{
NDIterCoalescing.FlipNegativeStrides(ref *_state);
}
if (_state->NDim > 1)
{
// NumPy's coalescing strategy depends on the order parameter:
//
// Key insight: Coalescing produces MEMORY-order iteration. This is correct for:
// - K-order: Memory order is exactly what we want
// - C-order on C-contiguous: Memory order == C-order
// - F-order on F-contiguous: Memory order == F-order
//
// But for F-order on C-contiguous (or C-order on F-contiguous), coalescing
// would produce the WRONG iteration order, so we must not coalesce.
//
// NumPy's behavior:
// - K-order on contiguous: Sort by stride, coalesce → memory order
// - K-order on non-contiguous: Fall back to C-order (no stride sorting)
// - C-order on C-contiguous: Sort by stride, coalesce → memory order (== C-order)
// - C-order on non-C-contiguous: Keep C-order, no coalescing
// - F-order on F-contiguous: Sort by stride, coalesce → memory order (== F-order)
// - F-order on C-contiguous: NO coalescing, reverse axes, iterate F-order
// Check contiguity once, use for all order decisions.
// allowFlip=true (absolute strides) only when FlipNegativeStrides will run,
// which is only for K-order (non-forced order).
// For C/F/A forced orders, negative strides are not contiguous since we
// preserve logical iteration order instead of memory order.
bool allowFlip = !isForcedOrder && (flags & NDIterGlobalFlags.DONT_NEGATE_STRIDES) == 0;
bool isCContiguous = CheckAllOperandsContiguous(true, allowFlip);
bool isFContiguous = CheckAllOperandsContiguous(false, allowFlip);
bool hasBroadcast = HasBroadcastStrides();
// For coalescing to work correctly:
// 1. All operands must be contiguous (either C or F order)
// - This includes reversed arrays (negative strides become positive after flip)
// 2. No broadcast dimensions (stride=0) - breaks stride-based sorting
bool isContiguous = (isCContiguous || isFContiguous) && !hasBroadcast;
// Determine effective order for non-contiguous arrays.
//
// NumPy K-order reorders axes by |stride| to match memory traversal even for
// non-contiguous views (e.g., transposed arrays). The only case where the
// stride-based sort produces wrong results is with BROADCAST axes (stride=0),
// because stride=0 breaks the ordering signal — we can't tell which broadcast
// axis should be innermost.
//
// So: fall back to C-order only when broadcast is present. For merely
// non-contiguous (transposed, strided views, negative strides), K-order does
// a proper descending-stride sort to match NumPy memory-order iteration.
NPY_ORDER effectiveOrder = order;
if ((order == NPY_ORDER.NPY_KEEPORDER || order == NPY_ORDER.NPY_ANYORDER) && hasBroadcast)
{
effectiveOrder = NPY_ORDER.NPY_CORDER;
}
if (!hasMultiIndex && !hasFlatIndex)
{
// Coalescing is possible when:
// - Arrays are contiguous in the REQUESTED order
// - No broadcast dimensions that would break stride-based sorting
// - No index tracking (C_INDEX/F_INDEX need original axis structure)
// Example: F-order on C-contiguous array should NOT coalesce
// (coalescing produces memory-order which is C-order, wrong for F-order)
bool canCoalesce;
if (order == NPY_ORDER.NPY_KEEPORDER || order == NPY_ORDER.NPY_ANYORDER)
{
// K-order: coalesce if contiguous in either C or F order
canCoalesce = isContiguous;
}
else if (order == NPY_ORDER.NPY_CORDER)
{
// C-order: coalesce only if C-contiguous (no broadcast)
canCoalesce = isCContiguous && !hasBroadcast;
}
else // NPY_FORTRANORDER
{
// F-order: coalesce only if F-contiguous (no broadcast)
canCoalesce = isFContiguous && !hasBroadcast;
}
if (canCoalesce)
{
// Sort axes by stride, then coalesce
NDIterCoalescing.ReorderAxesForCoalescing(ref *_state, NPY_ORDER.NPY_KEEPORDER, forCoalescing: true);
NDIterCoalescing.CoalesceAxes(ref *_state);
}
else
{
// Can't coalesce - reorder for the requested iteration order
NDIterCoalescing.ReorderAxesForCoalescing(ref *_state, effectiveOrder, forCoalescing: false);
// NumPy coalesces UNCONDITIONALLY after order resolution
// (npyiter_coalesce_axes has no contiguity guard): every
// size-1 axis is absorbed via the strict trivial branch AND
// every adjacent pair that all operands walk contiguously
// is merged — a C-contiguous (R, w) operand next to a 0-d
// scalar (strides 0, 0) collapses to ONE 1-D chunk, a
// trailing-narrow x[:, :, :w] to (R, w), a uniformly strided
// a[:, ::2] to one strided 1-D chunk. CoalesceAxes cannot
// do this here (it assumes the ascending pre-coalesce
// layout its own sort produces), so this branch used to
// only strip the size-1 axes and left every such case
// paying one kernel call per row. CoalesceAxesIterationOrder
// is the iteration-order merge (innermost axis last) and
// subsumes RemoveUnitAxes. Safe on this branch only: no
// multi-index/flat-index is tracked.
NDIterCoalescing.CoalesceAxesIterationOrder(ref *_state);
}
}
else
{
// With MULTI_INDEX or HASINDEX (C_INDEX/F_INDEX), just reorder axes
// without coalescing. Use effectiveOrder which applies K-order → C-order
// fallback for non-contiguous arrays.
NDIterCoalescing.ReorderAxesForCoalescing(ref *_state, effectiveOrder, forCoalescing: false);
}
}
// Set external loop flag separately (after coalescing)
if ((flags & NDIterGlobalFlags.EXTERNAL_LOOP) != 0)
{
_state->ItFlags |= (uint)NDIterFlags.EXLOOP;
}
// Set GROWINNER flag to maximize inner loop size during buffering
if ((flags & NDIterGlobalFlags.GROWINNER) != 0)
{
_state->ItFlags |= (uint)NDIterFlags.GROWINNER;
}
// Track multi-index if requested
if ((flags & NDIterGlobalFlags.MULTI_INDEX) != 0)
{
_state->ItFlags |= (uint)NDIterFlags.HASMULTIINDEX;
}
// Track flat index if requested (C_INDEX or F_INDEX)
if ((flags & NDIterGlobalFlags.C_INDEX) != 0)
{
_state->ItFlags |= (uint)NDIterFlags.HASINDEX;
_state->IsCIndex = true;
}
else if ((flags & NDIterGlobalFlags.F_INDEX) != 0)
{
_state->ItFlags |= (uint)NDIterFlags.HASINDEX;
_state->IsCIndex = false;
}
// Compute initial FlatIndex based on current coordinates (handles NEGPERM)
// Must be called after HASINDEX is set and negative strides are flipped
_state->InitializeFlatIndex();
// Update inner strides cache
// Note: CoalesceAxes calls this internally, but we need to ensure it's
// called even when coalescing is skipped (NDim <= 1 or MULTI_INDEX set)
if (_state->NDim <= 1 || (flags & NDIterGlobalFlags.MULTI_INDEX) != 0)
{
_state->UpdateInnerStrides();
}
// Update contiguity flags
UpdateContiguityFlags();
// Set up buffering if requested
if ((flags & NDIterGlobalFlags.BUFFERED) != 0)
{
_state->ItFlags |= (uint)NDIterFlags.BUFFER;
// An explicit buffersize wins; otherwise use the per-thread ufunc default
// (np.getbufsize()/np.setbufsize()), which is 8192 unless the caller changed it.
_state->BufferSize = bufferSize > 0 ? bufferSize : NDIterBufferManager.CurrentBufferSize;
bool isReduce = (_state->ItFlags & (uint)NDIterFlags.REDUCE) != 0;
if (!isReduce && (flags & NDIterGlobalFlags.DELAY_BUFALLOC) != 0)
{
// DELAY_BUFALLOC: defer allocation + first fill until first
// use. NumPy demands an explicit Reset before iterating
// (ValueError otherwise); NumSharp auto-ensures at the
// execution entry points (EnsureBuffersReady) instead, so
// construction stays allocation-free for the common
// construct-then-execute pattern. Ignored for buffered
// REDUCE, whose double-loop setup is construction-bound.
_state->ItFlags |= (uint)NDIterFlags.DELAYBUF;
}
else if (!isReduce)
{
// Windowed buffered iteration (non-reduce): allocate and
// prime the first window. FillBufferWindow swaps buffered
// operands' DataPtrs to their buffers and records the
// array write-back positions.
NDIterBufferManager.AllocateBuffers(ref *_state, _state->BufferSize);
NDIterBufferManager.FillBufferWindow(ref *_state,
NDIterBufferManager.ComputeTransferSize(ref *_state));
}
else
{
// Buffered REDUCE keeps its dedicated double-loop machinery
// (BufferedReduceIternext) and legacy initial fill. NOTE:
// that machinery predates masked copy-back; construction
// with WRITEMASKED succeeds (NumPy parity — 2.4.2 accepts
// the aligned-mask pattern) but the legacy write-back
// refuses loudly in CopyReduceBuffersToArrays instead of
// silently writing unmasked slots. Proper masked support
// lands with reductions-through-core (roadmap Wave 5).
NDIterBufferManager.AllocateBuffers(ref *_state, _state->BufferSize);
// Copy initial data to buffers (with casting if needed)
long copyCount = Math.Min(_state->IterSize, _state->BufferSize);
for (int op1 = 0; op1 < nop; op1++)
{
var opFlag = _state->GetOpFlags(op1);
if ((opFlag & NDIterOpFlags.READ) != 0 || (opFlag & NDIterOpFlags.READWRITE) != 0)
{
NDIterBufferManager.CopyToBuffer(ref *_state, op1, copyCount);
}
}
_state->BufIterEnd = copyCount;
SetupBufferedReduction(copyCount);
}
}
// Handle single iteration optimization
if (_state->IterSize <= 1)
{
_state->ItFlags |= (uint)NDIterFlags.ONEITERATION;
}
// PARALLEL_SAFE (NumSharp extension): the iteration range can be
// split across workers (RANGED + Copy() per worker) without write
// hazards when:
// - there is no REDUCE operand (cross-iteration accumulation on a
// shared slot; stretched write dims were flagged REDUCE above), and
// - there are no WRITE operands at all, or exactly one WRITE operand
// whose potential overlap with the inputs was already resolved by
// COPY_IF_OVERLAP processing (forced copy + write-back on Dispose).
// Costs nothing extra at construction — it only reuses decisions the
// overlap machinery already made. Consumed by the parallel ForEach
// work (roadmap Wave 6.2).
if ((_state->ItFlags & (uint)NDIterFlags.REDUCE) == 0)
{
int writeCount = 0;
for (int iop = 0; iop < nop; iop++)
{
if ((_state->GetOpFlags(iop) & NDIterOpFlags.WRITE) != 0)
writeCount++;
}
if (writeCount == 0 ||
(writeCount == 1 && (flags & NDIterGlobalFlags.COPY_IF_OVERLAP) != 0))
{
_state->ItFlags |= (uint)NDIterFlags.PARALLEL_SAFE;
}
}
}
/// <summary>
/// Compute iteration (broadcast) shape from operands.
/// Uses production NumSharp.Shape.ResolveReturnShape for standard broadcasting.
/// For op_axes, constructs a virtual shape per operand reflecting the mapping,
/// then broadcasts those virtual shapes together.
/// </summary>
[MethodImpl(MethodImplOptions.AggressiveOptimization)]
private static long[] CalculateBroadcastShape(int nop, NDArray[] op, NDIterPerOpFlags[] opFlags,
int opAxesNDim = -1, int[][]? opAxes = null)
{
// Validate null operands have ALLOCATE or VIRTUAL (NumPy
// npyiter_prepare_one_operand, nditer_constr.c:999-1007).
for (int i = 0; i < nop; i++)
{
if (op[i] is null &&
(opFlags[i] & (NDIterPerOpFlags.ALLOCATE | NDIterPerOpFlags.VIRTUAL)) == 0)
throw new ArgumentException(
"Iterator operand was NULL, but neither the " +
"ALLOCATE nor the VIRTUAL flag was specified");
}