@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
60 lines (59 loc) • 2.08 kB
TypeScript
import { DTypeGpu, GPUTensorConstructor, GPUTensorI } from '../../../tensor/gpu/interface';
import { GPUMemoryAllocator } from '../../../tensor/gpu/memory';
import { Input, Operation } from '../operation';
export interface GemmInfo {
shapeA?: readonly number[];
widthA?: number;
heightA?: number;
shapeB?: readonly number[];
widthB?: number;
heightB?: number;
shapeOutput?: readonly number[];
widthOutput?: number;
heightOutput?: number;
aTranspose?: boolean | number;
bTranspose?: boolean | number;
alpha?: number;
beta?: number;
M?: number;
N?: number;
O?: number;
rank?: number;
}
export interface GemmInput {
a: GPUTensorI;
b: GPUTensorI;
aTranspose: boolean;
bTranspose: boolean;
alpha: number;
beta: number;
}
export declare class GemmOperation<GPUTensor extends GPUTensorI, GemmInf extends GemmInfo = GemmInfo, GemmIn extends GemmInput = GemmInput> extends Operation<GPUTensor, GemmInf, GemmIn> {
protected maxIterations: number;
constructor(tensorConstructor: GPUTensorConstructor<GPUTensor>, dtype: DTypeGpu, allocator?: GPUMemoryAllocator);
getMainBody(): string;
getVariables(): string;
getFragmentShader(info: GemmInfo): string;
getTextureNames(): string[];
getUniformAttrs(): Input[];
calc(input: GemmInput): GPUTensor;
getOutputShape(input: GemmIn): readonly number[];
compile(info: GemmInf): void;
getCompilationInfo(input: GemmIn): GemmInf;
getInputInfoString(input: GemmIn): string;
}
export interface GemmCInfo extends GemmInfo {
shapeC?: readonly number[];
widthC?: number;
heightC?: number;
}
export interface GemmCInput extends GemmInput {
c: GPUTensorI;
}
export declare class GemmCOperation<GPUTensor extends GPUTensorI> extends GemmOperation<GPUTensor, GemmCInfo, GemmCInput> {
getTextureNames(): string[];
getFragmentShader(info: GemmInfo): string;
calc(input: GemmCInput): GPUTensor;
getCompilationInfo(input: GemmCInput): GemmCInfo;
getInputInfoString(input: GemmCInput): string;
}