@@ -1323,6 +1323,16 @@ export class Instance implements Disposable {
13231323 artifactCache : ArtifactCacheTemplate ,
13241324 signal ?: AbortSignal ,
13251325 ) {
1326+ const maxChunkBytes = 128 * 1024 * 1024 ;
1327+ const storageBytes = ( dtype : string ) => {
1328+ const match = dtype . match ( / ( \d + ) (?: x ( \d + ) ) ? $ / ) ;
1329+ if ( match === null ) {
1330+ throw new Error ( "Cannot determine storage width of dtype " + dtype ) ;
1331+ }
1332+ const bits = Number ( match [ 1 ] ) ;
1333+ const lanes = match [ 2 ] === undefined ? 1 : Number ( match [ 2 ] ) ;
1334+ return ( bits * lanes + 7 ) >> 3 ;
1335+ } ;
13261336 const perf = compact . getPerformance ( ) ;
13271337 const tstart = perf . now ( ) ;
13281338 let totalBytes = 0 ;
@@ -1421,9 +1431,62 @@ export class Instance implements Disposable {
14211431 this . empty ( rec . shape , rec . dtype , this . cpu ( ) )
14221432 )
14231433 } ) ;
1424- const recSource = buffer . slice ( rec . byteOffset , rec . byteOffset + rec . nbytes ) ;
1434+ const shardBytes = buffer instanceof Uint8Array ? buffer : new Uint8Array ( buffer ) ;
1435+ const recSource =
1436+ rec . byteOffset === 0 && rec . nbytes === shardBytes . byteLength
1437+ ? shardBytes
1438+ : shardBytes . subarray ( rec . byteOffset , rec . byteOffset + rec . nbytes ) ;
1439+ const canChunkRecord =
1440+ rec . nbytes > maxChunkBytes &&
1441+ rec . shape . length >= 1 &&
1442+ Number . isInteger ( rec . shape [ 0 ] ) &&
1443+ rec . shape [ 0 ] > 0 &&
1444+ rec . nbytes % rec . shape [ 0 ] === 0 ;
1445+ const outerDim = canChunkRecord ? rec . shape [ 0 ] : 1 ;
1446+ const sourceStrideBytes = canChunkRecord ? rec . nbytes / outerDim : rec . nbytes ;
1447+ const targetBytes = rec . shape . reduce ( ( acc , value ) => acc * value , 1 ) *
1448+ storageBytes ( rec . dtype ) ;
1449+ const targetStrideBytes = canChunkRecord ? targetBytes / outerDim : targetBytes ;
1450+ const copyRecordToTensor = ( targetTensor : Tensor , sourceBytes : Uint8Array ) => {
1451+ if ( ! canChunkRecord ) {
1452+ this . ctx . arrayDecodeStorage ( targetTensor , sourceBytes , rec . format , rec . dtype ) ;
1453+ return ;
1454+ }
1455+ const chunkOuterDim = Math . max ( 1 , Math . floor ( maxChunkBytes / sourceStrideBytes ) ) ;
1456+ for ( let outerOffset = 0 ; outerOffset < outerDim ; outerOffset += chunkOuterDim ) {
1457+ const outerCount = Math . min ( chunkOuterDim , outerDim - outerOffset ) ;
1458+ const sourceByteOffset = outerOffset * sourceStrideBytes ;
1459+ const targetByteOffset = outerOffset * targetStrideBytes ;
1460+ const chunkBytes = outerCount * sourceStrideBytes ;
1461+ const chunkShape = rec . shape . slice ( ) ;
1462+ chunkShape [ 0 ] = outerCount ;
1463+ // Wrap in withNewScope so TVM intermediate objects (shape tuple)
1464+ // are disposed after each chunk, but detach the view we need.
1465+ const chunkView = this . withNewScope ( ( ) => {
1466+ return this . detachFromCurrentScope (
1467+ this . ctx . tensorCreateView (
1468+ targetTensor ,
1469+ this . ctx . makeShapeTuple (
1470+ ...chunkShape . map ( ( value ) => new Scalar ( value , "int" ) ) ,
1471+ ) ,
1472+ rec . dtype ,
1473+ new Scalar ( targetByteOffset , "int" ) ,
1474+ )
1475+ ) ;
1476+ } ) ;
1477+ const chunkSource = sourceBytes . subarray (
1478+ sourceByteOffset ,
1479+ sourceByteOffset + chunkBytes ,
1480+ ) ;
1481+ try {
1482+ this . ctx . arrayDecodeStorage ( chunkView , chunkSource , rec . format , rec . dtype ) ;
1483+ } finally {
1484+ chunkView . dispose ( ) ;
1485+ }
1486+ }
1487+ } ;
14251488 // first sync copy to cpu.
1426- this . ctx . arrayDecodeStorage ( cpu_arr , new Uint8Array ( recSource ) , rec . format , rec . dtype ) ;
1489+ copyRecordToTensor ( cpu_arr , recSource ) ;
14271490 // then async stream into GPU if needed
14281491 if ( device . deviceType === DeviceStrToEnum . cpu ) {
14291492 this . tensorCacheUpdate ( rec . name , cpu_arr , false ) ;
@@ -1435,7 +1498,48 @@ export class Instance implements Disposable {
14351498 this . empty ( rec . shape , rec . dtype , device )
14361499 )
14371500 } ) ;
1438- gpu_arr . copyFrom ( cpu_arr ) ;
1501+ if ( ! canChunkRecord ) {
1502+ gpu_arr . copyFrom ( cpu_arr ) ;
1503+ } else {
1504+ const chunkOuterDim = Math . max ( 1 , Math . floor ( maxChunkBytes / sourceStrideBytes ) ) ;
1505+ for ( let outerOffset = 0 ; outerOffset < outerDim ; outerOffset += chunkOuterDim ) {
1506+ const outerCount = Math . min ( chunkOuterDim , outerDim - outerOffset ) ;
1507+ const targetByteOffset = outerOffset * targetStrideBytes ;
1508+ const chunkShape = rec . shape . slice ( ) ;
1509+ chunkShape [ 0 ] = outerCount ;
1510+ // Use withNewScope so the shape tuple is auto-disposed,
1511+ // and detach the views we need for manual lifetime control.
1512+ const [ cpuView , gpuView ] = this . withNewScope ( ( ) => {
1513+ const chunkShapeTuple = this . ctx . makeShapeTuple (
1514+ ...chunkShape . map ( ( value ) => new Scalar ( value , "int" ) ) ,
1515+ ) ;
1516+ return [
1517+ this . detachFromCurrentScope (
1518+ this . ctx . tensorCreateView (
1519+ cpu_arr ,
1520+ chunkShapeTuple ,
1521+ rec . dtype ,
1522+ new Scalar ( targetByteOffset , "int" ) ,
1523+ )
1524+ ) ,
1525+ this . detachFromCurrentScope (
1526+ this . ctx . tensorCreateView (
1527+ gpu_arr ,
1528+ chunkShapeTuple ,
1529+ rec . dtype ,
1530+ new Scalar ( targetByteOffset , "int" ) ,
1531+ )
1532+ ) ,
1533+ ] ;
1534+ } ) ;
1535+ try {
1536+ gpuView . copyFrom ( cpuView ) ;
1537+ } finally {
1538+ cpuView . dispose ( ) ;
1539+ gpuView . dispose ( ) ;
1540+ }
1541+ }
1542+ }
14391543 await device . sync ( ) ;
14401544 this . tensorCacheUpdate ( rec . name , gpu_arr , false ) ;
14411545 cpu_arr . dispose ( ) ;
@@ -2258,6 +2362,25 @@ export class Instance implements Disposable {
22582362 case TypeIndex . kTVMFFIOpaquePtr : {
22592363 return this . memory . loadPointer ( valuePtr ) ;
22602364 }
2365+ case TypeIndex . kTVMFFIShape : {
2366+ const shapeObjPtr = this . memory . loadPointer ( valuePtr ) ;
2367+ if ( callbackArg ) {
2368+ const shapeCellPtr = shapeObjPtr + SizeOf . ObjectHeader ;
2369+ const shapeDataPtr = this . memory . loadPointer ( shapeCellPtr ) ;
2370+ const shapeLen = this . memory . loadUSize ( shapeCellPtr + this . memory . sizeofPtr ( ) ) ;
2371+ const result = new Array < number > ( shapeLen ) ;
2372+ for ( let i = 0 ; i < shapeLen ; ++ i ) {
2373+ result [ i ] = this . memory . loadI64 ( shapeDataPtr + i * SizeOf . I64 ) ;
2374+ }
2375+ this . lib . checkCall (
2376+ ( this . lib . exports . TVMFFIObjectDecRef as ctypes . FTVMFFIObjectDecRef ) ( shapeObjPtr )
2377+ ) ;
2378+ return result ;
2379+ }
2380+ return this . ctx . attachToCurrentScope (
2381+ new TVMObject ( shapeObjPtr , this . lib , this . ctx )
2382+ ) ;
2383+ }
22612384 case TypeIndex . kTVMFFITensor : {
22622385 return this . ctx . attachToCurrentScope (
22632386 new Tensor ( this . memory . loadPointer ( valuePtr ) , this . lib , this . ctx , false )
0 commit comments