gx
chenyc
2025-06-12 7b72ac13a83764a662159d4a49b7fffb90476ecb
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
/**
 * @license
 * Copyright 2017 Google LLC. All Rights Reserved.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * =============================================================================
 */
/// <amd-module name="@tensorflow/tfjs-backend-webgl/dist/gpgpu_math" />
import { Tensor, TypedArray } from '@tensorflow/tfjs-core';
import { GPGPUContext, GPGPUContextProgram } from './gpgpu_context';
import { ShapeInfo, UniformType } from './shader_compiler';
import { PackingScheme, TextureData, TextureUsage } from './tex_util';
export interface GPGPUProgram {
    variableNames: string[];
    outputShape: number[];
    userCode: string;
    enableShapeUniforms?: boolean;
    /** If true, this program expects packed input textures. Defaults to false. */
    packedInputs?: boolean;
    /** If true, this program produces a packed texture. Defaults to false. */
    packedOutput?: boolean;
    /**
     * Affects what type of texture we allocate for the output. Defaults to
     * `TextureUsage.RENDER`.
     */
    outTexUsage?: TextureUsage;
    /**
     * The type of scheme to use when packing texels for the output values.
     * See `PackingScheme` for details. Defaults to `PackingScheme.SHARED_BATCH`.
     */
    outPackingScheme?: PackingScheme;
    customUniforms?: Array<{
        name: string;
        arrayIndex?: number;
        type: UniformType;
    }>;
}
export interface GPGPUBinary extends GPGPUBinaryLocations {
    webGLProgram: GPGPUContextProgram;
    program: GPGPUProgram;
    source: string;
    fragmentShader: WebGLShader;
    inShapeInfos: ShapeInfo[];
    outShapeInfo: ShapeInfo;
}
export interface GPGPUBinaryLocations {
    customUniformLocations?: WebGLUniformLocation[];
    infLoc: WebGLUniformLocation;
    nanLoc: WebGLUniformLocation;
    outShapeLocation?: WebGLUniformLocation;
    outShapeStridesLocation?: WebGLUniformLocation;
    outTexShapeLocation?: WebGLUniformLocation;
    variablesLocations?: GPGPUVariableLocations[];
}
export interface GPGPUVariableLocations {
    name: string;
    uniform: WebGLUniformLocation;
    offset: WebGLUniformLocation;
    shape?: WebGLUniformLocation;
    texShape?: WebGLUniformLocation;
}
export interface TensorData {
    shape: number[];
    texData: TextureData;
    isUniform: boolean;
    uniformValues?: TypedArray;
}
export declare function compileProgram<T extends Tensor, K extends Tensor>(gpgpu: GPGPUContext, program: GPGPUProgram, inputs: TensorData[], output: TensorData): GPGPUBinary;
export declare function getUniformLocations(gpgpu: GPGPUContext, program: GPGPUProgram, webGLProgram: WebGLProgram): GPGPUBinaryLocations;
export declare function runProgram<T extends Tensor, K extends Tensor>(gpgpu: GPGPUContext, binary: GPGPUBinary, inputs: TensorData[], output: TensorData, customUniformValues?: number[][]): void;
export declare function makeShaderKey(program: GPGPUProgram, inputs: TensorData[], output: TensorData): string;
export declare function useShapeUniforms(rank: number): boolean;