From 1d4a3f3600b227c52653ef14df7356eb076cc809 Mon Sep 17 00:00:00 2001 From: Christian Legnitto Date: Sat, 23 Nov 2024 14:39:11 -0400 Subject: [PATCH] Make snippets more robust, correct line numbers * Snippets can now shift code left while keeping other indentation. * If lines are specified, a hash needs to be specified. * Update line numbers so they are correct. --- .../2024-11-21-optimizing-matrix-mul/index.md | 11 ++- .../snippets/naive.tsx | 9 ++- .../snippets/party.tsx | 19 +++-- .../snippets/workgroup_256.tsx | 5 +- .../snippets/workgroup_2d.tsx | 32 +------- src/components/Snippet/index.tsx | 75 ++++++++++++++++--- 6 files changed, 97 insertions(+), 54 deletions(-) diff --git a/blog/2024-11-21-optimizing-matrix-mul/index.md b/blog/2024-11-21-optimizing-matrix-mul/index.md index 96235d4..71c99d6 100644 --- a/blog/2024-11-21-optimizing-matrix-mul/index.md +++ b/blog/2024-11-21-optimizing-matrix-mul/index.md @@ -128,9 +128,9 @@ import { WebGpuKernel } from './snippets/naive.tsx'; With Rust GPU, we specify the inputs as arguments to the kernel and configure them with [procedural macros](https://doc.Rust-lang.org/reference/procedural-macros.html): -import { RustNaiveInputs } from './snippets/naive.tsx'; +import { RustNaiveKernel } from './snippets/naive.tsx'; - + This code looks like normal Rust code but _runs entirely on the GPU._ @@ -301,6 +301,13 @@ improvement over the last kernel. To stay true to the spirit of Zach's original blog post, we'll wrap things up here and leave the "fancier" experiments for another time. +### A note on performance + +I didn't include performance numbers as I have a different machine than Zach. The +complete runnable code can be [found on +GitHub](https://github.com/Rust-GPU/rust-gpu.github.io/tree/main/blog/2024-11-21-optimizing-matrix-mul/code) +and you can run the benchmarks yourself with `cargo bench`. + ## Reflections on porting to Rust GPU Porting to Rust GPU went quickly, as the kernels Zach used were fairly simple. Most of diff --git a/blog/2024-11-21-optimizing-matrix-mul/snippets/naive.tsx b/blog/2024-11-21-optimizing-matrix-mul/snippets/naive.tsx index d637a49..e5232c6 100644 --- a/blog/2024-11-21-optimizing-matrix-mul/snippets/naive.tsx +++ b/blog/2024-11-21-optimizing-matrix-mul/snippets/naive.tsx @@ -42,11 +42,10 @@ fn main(@builtin(global_invocation_id) global_id: vec3) { ); -export const RustNaiveInputs: React.FC = () => ( +export const RustNaiveKernel: React.FC = () => ( @@ -59,6 +58,7 @@ export const RustNaiveWorkgroupCount: React.FC = () => ( language="rust" className="text-xs" lines="26-34" + hash="8abb43d" title="Calculating on the CPU how many workgroup dispatches are needed" > {RustWorkgroupCount} @@ -69,7 +69,8 @@ export const RustNaiveDispatch: React.FC = () => ( @@ -78,7 +79,7 @@ export const RustNaiveDispatch: React.FC = () => ( ); export const RustNaiveWorkgroup: React.FC = () => ( - + {RustKernelSource} ); diff --git a/blog/2024-11-21-optimizing-matrix-mul/snippets/party.tsx b/blog/2024-11-21-optimizing-matrix-mul/snippets/party.tsx index 0f3b99f..d16d930 100644 --- a/blog/2024-11-21-optimizing-matrix-mul/snippets/party.tsx +++ b/blog/2024-11-21-optimizing-matrix-mul/snippets/party.tsx @@ -7,7 +7,7 @@ import RustWgpuBackend from "!!raw-loader!../code/crates/cpu/matmul/src/backends import RustCpuBackendSource from "!!raw-loader!../code/crates/cpu/matmul/src/backends/cpu.rs"; export const RustPartySettings: React.FC = () => ( - + {RustKernelSource} ); @@ -19,13 +19,19 @@ export const RustIsomorphic: React.FC = () => ( ); export const RustIsomorphicGlam: React.FC = () => ( - + {RustIsomorphicSource} ); export const RustIsomorphicDeps: React.FC = () => ( - + {RustIsomorphicCargoToml} ); @@ -33,7 +39,8 @@ export const RustIsomorphicDeps: React.FC = () => ( export const RustWgpuDimensions: React.FC = () => ( @@ -42,13 +49,13 @@ export const RustWgpuDimensions: React.FC = () => ( ); export const RustCpuBackendHarness: React.FC = () => ( - + {RustCpuBackendSource} ); export const RustCpuBackendTest: React.FC = () => ( - + {RustCpuBackendSource} ); diff --git a/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_256.tsx b/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_256.tsx index def6ba2..506dac6 100644 --- a/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_256.tsx +++ b/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_256.tsx @@ -4,7 +4,7 @@ import RustKernelSource from "!!raw-loader!../code/crates/gpu/workgroup_256/src/ import VariantsSource from "!!raw-loader!../code/crates/cpu/matmul/src/variants.rs"; export const RustWorkgroup256Workgroup: React.FC = () => ( - + {RustKernelSource} ); @@ -13,7 +13,8 @@ export const RustWorkgroup256WorkgroupCount: React.FC = () => ( {VariantsSource} diff --git a/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_2d.tsx b/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_2d.tsx index 39ebf10..e8df4ea 100644 --- a/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_2d.tsx +++ b/blog/2024-11-21-optimizing-matrix-mul/snippets/workgroup_2d.tsx @@ -15,44 +15,14 @@ export const RustWorkgroup2d: React.FC = () => ( ); -/* -export const RustWorkgroup2d: React.FC = () => ( - - {RustKernelSource} - -); -*/ - -export const RustWorkgroup2dWorkgroup: React.FC = () => ( - - {RustKernelSource} - -); - export const RustWorkgroup2dWorkgroupCount: React.FC = () => ( {VariantsSource} ); - -export const RustWorkgroup2dWgpuDispatch: React.FC = () => ( - - {WgpuBackendSource} - -); diff --git a/src/components/Snippet/index.tsx b/src/components/Snippet/index.tsx index 3cc64a7..4551a89 100644 --- a/src/components/Snippet/index.tsx +++ b/src/components/Snippet/index.tsx @@ -1,4 +1,4 @@ -import React from "react"; +import React, { useEffect, useState } from "react"; import CodeBlock from "@theme/CodeBlock"; interface SnippetProps extends React.ComponentProps { @@ -8,27 +8,73 @@ interface SnippetProps extends React.ComponentProps { lines?: string; omitted_placeholder?: string; strip_leading_spaces?: boolean; + /** + * Optional short hash of the content (first N characters of SHA-256), + * required only when `lines` is specified. + */ + hash?: string; } /** * A component for rendering a snippet of code, optionally filtering lines, - * showing ellipses for omissions, and stripping all leading spaces. + * showing ellipses for omissions, stripping leading spaces, and validating hash. */ const Snippet: React.FC = ({ children, lines, omitted_placeholder = "...", strip_leading_spaces = false, + hash, ...props }) => { + const [error, setError] = useState(null); + if (typeof children !== "string") { - console.error( + throw new Error( "Snippet expects children to be a string containing the file content." ); - return null; } - // Parse the `linesToInclude` metadata string into an array of line numbers. + /** + * Utility function to compute the SHA-256 hash of a string. + * @param content The input string + * @returns Promise resolving to a hex-encoded hash + */ + const computeHash = async (content: string): Promise => { + const encoder = new TextEncoder(); + const data = encoder.encode(content); + const hashBuffer = await crypto.subtle.digest("SHA-256", data); + return Array.from(new Uint8Array(hashBuffer)) + .map((byte) => byte.toString(16).padStart(2, "0")) + .join(""); + }; + + useEffect(() => { + if (lines) { + computeHash(children).then((computedHash) => { + const shortHash = computedHash.slice(0, 7); // Use 7 characters for the short hash + + if (!hash) { + setError( + `The \`hash\` prop is required when \`lines\` is specified.\n` + + `Provide the following hash as the \`hash\` prop: ${shortHash}` + ); + } else if (shortHash !== hash) { + setError( + `Snippet hash mismatch.\n` + + `Specified: ${hash}, but content is: ${shortHash} (full hash: ${computedHash}).\n` + + `Check if the line numbers are still relevant and update the hash.` + ); + } + }); + } + }, [children, lines, hash]); + + if (error) { + throw new Error(error); + } + + // Parse the `lines` metadata string into an array of line numbers. const parseLineRanges = (metaString?: string): number[] => { if (!metaString) return []; return metaString.split(",").flatMap((range) => { @@ -46,16 +92,27 @@ const Snippet: React.FC = ({ if (lines.length === 0) return content; // If no specific lines are specified, return full content. const includedContent: string[] = []; + + // Filter lines and find the minimum indentation + const selectedLines = lines + .map((line) => allLines[line - 1] || "") + .filter((line) => line.trim().length > 0); // Ignore blank lines + + const minIndent = selectedLines.reduce((min, line) => { + const indentMatch = line.match(/^(\s*)\S/); + const indentLength = indentMatch ? indentMatch[1].length : 0; + return Math.min(min, indentLength); + }, Infinity); + lines.forEach((line, index) => { if (index > 0 && lines[index - 1] < line - 1) { includedContent.push(omitted_placeholder); // Add placeholder for omitted lines } const rawLine = allLines[line - 1] || ""; - const formattedLine = strip_leading_spaces - ? rawLine.trimStart() - : rawLine; - includedContent.push(formattedLine); + const trimmedLine = + rawLine.trim().length > 0 ? rawLine.slice(minIndent) : rawLine; + includedContent.push(trimmedLine); }); // Add placeholder if lines at the end are omitted