Skip to content

fix: Remove 'use gpu' after transformation #2706

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aleksanderkatan merged 8 commits into from
Jul 10, 2026
Merged

fix: Remove 'use gpu' after transformation #2706

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
2f74d79
Add tests
aleksanderkatan Jul 9, 2026
07e2bab
Add method stubs
aleksanderkatan Jul 9, 2026
0f955bd
remove directive in rollup
aleksanderkatan Jul 9, 2026
c03764a
Remove directive in babel
aleksanderkatan Jul 10, 2026
5cdb859
Update tests
aleksanderkatan Jul 10, 2026
7fcb443
Remove `alreadyTransformed`, add new test
aleksanderkatan Jul 10, 2026
6e50b82
Merge branch 'main' into fix/remove-use-gpu
aleksanderkatan Jul 10, 2026
985409a
Merge branch 'main' into fix/remove-use-gpu
aleksanderkatan Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/typegpu/tests/internal/autoname.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ describe('autonaming', () => {
expect(scope.toString()).toMatchInlineSnapshot(`
"() => {
const myFun = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function myFun() {
"use gpu";

return 0;
}), {
v: 1,
Expand Down
9 changes: 8 additions & 1 deletion packages/unplugin-typegpu/src/babel.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { transpileFn } from 'tinyest-for-wgsl';
import * as t from '@babel/types';
import {
METADATA_FORMAT_VERSION,
type MetadatableFunction,
type PluginState,
defaultOptions,
functionVisitor,
Expand Down Expand Up @@ -101,7 +102,6 @@ function assignMetadata(
if (visibility) {
// Hoisting the declaration to the top of the scope
visibility.unshiftContainer('body', replacement as t.Statement);
this.alreadyTransformed.add(expression);

const id = t.isFunctionDeclaration(path.node) ? path.node.id : undefined;
if (id && path.parentPath.isExportNamedDeclaration()) {
Expand Down Expand Up @@ -148,6 +148,12 @@ function replaceWithAssignmentOverload(
);
}

function removeUseGpuDirective(this: PluginState, path: NodePath<MetadatableFunction>) {
const directives = path.get('body').get('directives');
const maybeUseGpu = directives.find((directive) => directive.node.value.value === 'use gpu');
maybeUseGpu?.remove();
}

function replaceWithBinaryOverload(path: NodePath<t.BinaryExpression>, runtimeFn: string): void {
path.replaceWith(
t.callExpression(i(runtimeFn), [path.node.left as t.Expression, path.node.right]),
Expand All @@ -165,6 +171,7 @@ export default function TypeGPUPlugin() {
wrapInAutoName,
replaceWithAssignmentOverload,
replaceWithBinaryOverload,
removeUseGpuDirective,
});
},
visitor: {
Expand Down
11 changes: 3 additions & 8 deletions packages/unplugin-typegpu/src/core/common.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,8 @@ export interface TransformMethods {
replaceWithAssignmentOverload(path: NodePath<t.AssignmentExpression>, runtimeFn: string): void;

replaceWithBinaryOverload(path: NodePath<t.BinaryExpression>, runtimeFn: string): void;

removeUseGpuDirective(this: PluginState, path: NodePath<MetadatableFunction>): void;
}

export interface PluginState extends TransformMethods {
Expand All @@ -107,8 +109,6 @@ export interface PluginState extends TransformMethods {
opts: Options;

inUseGpuScope: boolean;

alreadyTransformed: WeakSet<t.Node>;
}

export interface NodeLocation {
Expand All @@ -127,7 +127,6 @@ export function initPluginState(state: PluginState, methods: TransformMethods):
state.tgpuAliases = new Set<string>(state.opts.forceTgpuAlias ? [state.opts.forceTgpuAlias] : []);
state.autoNamingEnabled = state.opts.autoNamingEnabled ?? true;
state.inUseGpuScope = false;
state.alreadyTransformed = new WeakSet<t.Node>();
Object.assign(state, methods);
}

Expand Down Expand Up @@ -450,20 +449,16 @@ function functionOnExit(
if (!containsUseGpuDirective(node)) {
return;
}
state.removeUseGpuDirective(path);

state.inUseGpuScope = false;

if (state.alreadyTransformed.has(node)) {
return;
}

const ast = fnNodeToTranspiledMap.get(path.node);
const maybeName = getFunctionName(path);
if (!ast) {
throw new Error(`No metadata found for function ${maybeName ?? '<unnamed>'}`);
}
state.assignMetadata(path, maybeName, ast);
state.alreadyTransformed.add(node);
path.skip();
}

Expand Down
10 changes: 10 additions & 0 deletions packages/unplugin-typegpu/src/core/factory.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,15 @@ function replaceWithBinaryOverload(
this.overwrite(path.node, `${runtimeFn}(${lhs}, ${rhs})`);
}

function removeUseGpuDirective(this: UnpluginPluginState, path: NodePath<MetadatableFunction>) {
const directives = 'directives' in path.node.body ? (path.node.body?.directives ?? []) : [];
for (const directive of directives) {
if (directive.value.value === 'use gpu') {
this.remove(directive);
}
}
}

const NodeUtils = {
slice(this: UnpluginPluginState, node: NodeLocation): string {
return this.magicString.slice(node.start ?? 0, node.end ?? 0);
Expand Down Expand Up @@ -188,6 +197,7 @@ export const unpluginFactory = ((rawOptions, _meta) => {
wrapInAutoName,
replaceWithAssignmentOverload,
replaceWithBinaryOverload,
removeUseGpuDirective,
});

traverse(ast, functionVisitor, undefined, state);
Expand Down
24 changes: 7 additions & 17 deletions packages/unplugin-typegpu/test/auto-naming.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -327,8 +327,6 @@ describe('[BABEL] auto naming', () => {

expect(babelTransform(code, { autoNamingEnabled: true })).toMatchInlineSnapshot(`
"const myFun3 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function myFun3() {
'use gpu';

return 0;
}, {
v: 1,
Expand All @@ -343,8 +341,6 @@ describe('[BABEL] auto naming', () => {
}
}) && $.f)({});
const myFun1 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {
'use gpu';

return 0;
}, {
v: 1,
Expand All @@ -359,8 +355,6 @@ describe('[BABEL] auto naming', () => {
}
}) && $.f)({});
const myFun2 = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function () {
'use gpu';

return 0;
}, {
v: 1,
Expand Down Expand Up @@ -510,9 +504,7 @@ describe('[BABEL] auto naming', () => {
expect(babelTransform(code, { autoNamingEnabled: true })).toMatchInlineSnapshot(`
"import { tgpu } from 'typegpu';
const root = await tgpu.init();
const myGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {
'use gpu';
}, {
const myGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, {
v: 1,
name: undefined,
ast: {
Expand All @@ -524,9 +516,7 @@ describe('[BABEL] auto naming', () => {
return {};
}
}) && $.f)({})), "myGuardedPipeline");
const anotherGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {
'use gpu';
}, {
const anotherGuardedPipeline = /*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline(/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = () => {}, {
v: 1,
name: undefined,
ast: {
Expand Down Expand Up @@ -829,7 +819,7 @@ describe('[ROLLUP] auto naming', () => {

expect(await rollupTransform(code, { autoNamingEnabled: true })).toMatchInlineSnapshot(`
"const myFun3 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function myFun3() {
'use gpu';

return 0;
}), {
v: 1,
Expand All @@ -839,7 +829,7 @@ describe('[ROLLUP] auto naming', () => {
}) && $.f)({}));

const myFun1 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
'use gpu';

return 0;
}), {
v: 1,
Expand All @@ -849,7 +839,7 @@ describe('[ROLLUP] auto naming', () => {
}) && $.f)({}));

const myFun2 = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function () {
'use gpu';

return 0;
}), {
v: 1,
Expand Down Expand Up @@ -1027,7 +1017,7 @@ describe('[ROLLUP] auto naming', () => {
const root = await tgpu.init();

const myGuardedPipeline = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root.createGuardedComputePipeline((/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
'use gpu';

}), {
v: 1,
name: undefined,
Expand All @@ -1037,7 +1027,7 @@ describe('[ROLLUP] auto naming', () => {

const anotherGuardedPipeline = (/*#__PURE__*/(globalThis.__TYPEGPU_AUTONAME__ ?? (a => a))(root
.createGuardedComputePipeline((/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (() => {
'use gpu';

}), {
v: 1,
name: undefined,
Expand Down
2 changes: 0 additions & 2 deletions packages/unplugin-typegpu/test/typescript-syntax.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,6 @@ describe('as type', () => {

expect(babelTransform(code)).toMatchInlineSnapshot(`
"const hello = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a: number, b: number | undefined) => {
'use gpu';

return __tsover_add(a, b as number);
}, {
v: 1,
Expand Down
Loading
Loading