This tool allows you to explore and visualize tensor operations.
+If you're new to Triton or need to brush up on your skills, check out Triton Puzzles. It's a collection of Triton problems designed to get you up to speed:
+ +These puzzles provide hands-on experience with Triton concepts, which can help you better understand the visualizations in this tool.
+ ` + }, + { + title: "Using the Sliders", + content: ` +The sliders on the right side of the screen allow you to filter the grid blocks:
+Set a slider to -1 to show all blocks along that dimension.
+ ` + }, + { + title: "Detailed View", + content: ` +Clicking on a grid block will open a detailed view of the operations for that block.
+In the detailed view, you can:
+Triton operates on slices of tensors rather than entire tensors. This is a key concept in understanding how Triton kernels process data efficiently.
+ +Inside Triton kernels, tl.load is used to load slices of the global tensor. For example:
a = tl.load(a_ptrs, mask=offs_k[None, :])This operation loads only the necessary slice of data, optimizing memory access and computation.
+ +tl.load and tl.store operations in Triton handle the slicing automatically.
+ Diagram showing how a global PyTorch tensor is sliced for processing in Triton kernels.
+ ` + }, + { + title: "Memory Hierarchy in AI Accelerators and Triton Operations", + content: ` +Understanding the memory hierarchy in AI accelerators is crucial for optimizing Triton kernels. While this example uses a GPU, it's important to note that Triton is designed for various AI accelerators, not just GPUs.
+ +Triton provides mechanisms to efficiently move data between these memory types across various accelerators:
+tl.load: Generally used to load a specified portion of the global tensor into shared memory*. This operation can significantly speed up subsequent computations.tl.store: Used to write processed data back to global memory.* It's important to note that Triton doesn't always transfer tensor slices to shared memory. This operation is performed only when it makes sense for performance optimization.
+ +tl.load and tl.store operations abstract the complexity of memory transfers, making code portable across different accelerator types.
+ Diagram showing the memory hierarchy of an RTX 4090 GPU as an example of AI accelerator architecture. Global memory (magenta) and on-chip memory including shared memory* (red) are highlighted. Note that this structure is similar across many AI accelerators.
+ ` + } +]; + +let currentPage = 0; + +export function createInfoPopup() { + const infoPopup = document.createElement('div'); + infoPopup.style.display = 'none'; + infoPopup.style.position = 'fixed'; + infoPopup.style.top = '5%'; + infoPopup.style.left = '5%'; + infoPopup.style.width = '85%'; + infoPopup.style.height = '75%'; + infoPopup.style.backgroundColor = '#1a1a1a'; + infoPopup.style.color = '#fff'; + infoPopup.style.padding = '30px'; + infoPopup.style.borderRadius = '15px'; + infoPopup.style.boxShadow = '0 4px 6px rgba(0, 0, 0, 0.5)'; + infoPopup.style.zIndex = '1000'; + infoPopup.style.overflow = 'auto'; + + const closeButton = createButton('×', () => { infoPopup.style.display = 'none'; }); + closeButton.style.position = 'absolute'; + closeButton.style.top = '20px'; + closeButton.style.right = '20px'; + closeButton.style.fontSize = '36px'; + closeButton.style.color = '#fff'; + + const content = document.createElement('div'); + content.id = 'info-content'; + content.style.fontSize = '18px'; + content.style.lineHeight = '1.6'; + + const navigation = document.createElement('div'); + navigation.style.display = 'flex'; + navigation.style.justifyContent = 'space-between'; + navigation.style.marginTop = '30px'; + + const prevButton = createNavButton('← Previous', () => navigatePage(-1)); + const nextButton = createNavButton('Next →', () => navigatePage(1)); + + navigation.appendChild(prevButton); + navigation.appendChild(nextButton); + + infoPopup.appendChild(closeButton); + infoPopup.appendChild(content); + infoPopup.appendChild(navigation); + + document.body.appendChild(infoPopup); + + updateContent(); + + return infoPopup; +} + +function createButton(text, onClick) { + const button = document.createElement('button'); + button.innerHTML = text; + button.style.border = 'none'; + button.style.background = 'none'; + button.style.cursor = 'pointer'; + button.style.color = '#fff'; + button.onclick = onClick; + return button; +} + +function createNavButton(text, onClick) { + const button = createButton(text, onClick); + button.style.fontSize = '20px'; + button.style.padding = '15px 25px'; + button.style.backgroundColor = '#333'; + button.style.borderRadius = '8px'; + button.style.transition = 'background-color 0.3s'; + button.onmouseover = () => { button.style.backgroundColor = '#444'; }; + button.onmouseout = () => { button.style.backgroundColor = '#333'; }; + return button; +} + +function navigatePage(direction) { + currentPage += direction; + if (currentPage < 0) currentPage = infoContent.length - 1; + if (currentPage >= infoContent.length) currentPage = 0; + updateContent(); +} + +function updateContent() { + const contentDiv = document.getElementById('info-content'); + const pageContent = infoContent[currentPage]; + contentDiv.innerHTML = ` +X: ${x + 1}
+Y: ${y + 1}
+Z: ${z + 1}
+Dimensions: ${dims.join(' x ')}
+Value: ${value !== undefined ? value : 'Loading...'}
+ `; + } + + function updateInfoPanel(globalCoord, sliceCoord, index) { + sideMenu.innerHTML = ` +Global Coords: (${globalCoord.join(', ')})
+Slice Coords: (${sliceCoord.join(', ')})
+Progress: ${index + 1}/${op.global_coords.length}
+ `; + } + + function createSideMenu(container) { + const menu = document.createElement('div'); + menu.style.position = 'absolute'; + menu.style.top = '10px'; + menu.style.right = '10px'; + menu.style.width = '200px'; + menu.style.padding = '10px'; + menu.style.backgroundColor = 'rgba(0, 0, 0, 0.7)'; + menu.style.color = 'white'; + menu.style.fontFamily = 'Arial, sans-serif'; + menu.style.fontSize = '14px'; + menu.style.borderRadius = '5px'; + container.appendChild(menu); + return menu; + } + + function addLabels(scene, globalTensor, sliceTensor) { + addLabel(scene, "Global Tensor", globalTensor.position); + addLabel(scene, "Slice Tensor", sliceTensor.position); + } + + function addLabel(scene, text, position) { + const canvas = document.createElement('canvas'); + const context = canvas.getContext('2d'); + context.font = 'Bold 24px Arial'; + context.fillStyle = 'white'; + context.fillText(text, 0, 24); + + const texture = new THREE.CanvasTexture(canvas); + const material = new THREE.SpriteMaterial({ map: texture }); + const sprite = new THREE.Sprite(material); + sprite.position.set(position.x, position.y + 2, position.z); + sprite.scale.set(4, 2, 1); + scene.add(sprite); + } + +} diff --git a/triton_viz/static/load_utils.js b/triton_viz/static/load_utils.js new file mode 100644 index 0000000..2729d04 --- /dev/null +++ b/triton_viz/static/load_utils.js @@ -0,0 +1,200 @@ +import * as THREE from 'https://cdn.jsdelivr.net/npm/three@0.155.0/build/three.module.js'; + +export const CUBE_SIZE = 0.2; +export const GAP = 0.05; +export const COLOR_HOVER = new THREE.Color(1.0, 1.0, 0.0); // Yellow for hover effect +export const COLOR_EDGE = new THREE.Color(0.5, 0.5, 0.5); // Gray (for cube edges) + +const COLOR_SLICE = new THREE.Color(0.0, 0.7, 1.0); // Cyan (starting color for global slice) + +export function setupScene(container, backgroundColor = 0x000000) { + const scene = new THREE.Scene(); + scene.background = new THREE.Color(backgroundColor); + const camera = new THREE.PerspectiveCamera(45, container.clientWidth / container.clientHeight, 0.1, 1000); + const renderer = new THREE.WebGLRenderer({ antialias: true }); + renderer.setSize(container.clientWidth, container.clientHeight); + container.appendChild(renderer.domElement); + + const ambientLight = new THREE.AmbientLight(0xffffff, 0.5); + scene.add(ambientLight); + const directionalLight = new THREE.DirectionalLight(0xffffff, 0.5); + directionalLight.position.set(10, 10, 10); + scene.add(directionalLight); + + return { scene, camera, renderer }; +} + +export function setupGeometries() { + const cubeGeometry = new THREE.BoxGeometry(CUBE_SIZE, CUBE_SIZE, CUBE_SIZE); + const edgesGeometry = new THREE.EdgesGeometry(cubeGeometry); + const lineMaterial = new THREE.LineBasicMaterial({ color: COLOR_EDGE }); + return { cubeGeometry, edgesGeometry, lineMaterial }; +} + +export function createCube(color, tensorName, x, y, z, cubeGeometry, edgesGeometry, lineMaterial) { + const cubeMaterial = new THREE.MeshPhongMaterial({ color: color }); + const cube = new THREE.Mesh(cubeGeometry, cubeMaterial); + const edges = new THREE.LineSegments(edgesGeometry, lineMaterial); + cube.add(edges); + + const hoverGeometry = new THREE.BoxGeometry(CUBE_SIZE * 1.05, CUBE_SIZE * 1.05, CUBE_SIZE * 1.05); + const hoverEdgesGeometry = new THREE.EdgesGeometry(hoverGeometry); + const hoverOutline = new THREE.LineSegments(hoverEdgesGeometry, new THREE.LineBasicMaterial({ color: COLOR_HOVER })); + hoverOutline.visible = false; + hoverOutline.name = 'hoverOutline'; + cube.add(hoverOutline); + + // Add custom properties to store tensor coordinates + cube.userData.tensor0 = z; + cube.userData.tensor1 = y; + cube.userData.tensor2 = x; + cube.userData.tensorName = tensorName; + + cube.name = `${tensorName}_cube_${x}_${y}_${z}`; + + return cube; +} + +export function createTensor(shape, coords, color, tensorName, cubeGeometry, edgesGeometry, lineMaterial) { + console.log(`Creating ${tensorName} tensor:`, shape, coords); + const tensor = new THREE.Group(); + let [width, height, depth] = shape; + depth = depth || 1; + height = height || 1; + + if (tensorName === 'Global') { + console.log(`Creating global tensor with dimensions: ${width}x${height}x${depth}`); + for (let z = 0; z < depth; z++) { + for (let y = 0; y < height; y++) { + for (let x = 0; x < width; x++) { + const cube = createCube(color, tensorName, x, y, z, cubeGeometry, edgesGeometry, lineMaterial); + cube.position.set( + x * (CUBE_SIZE + GAP), + -y * (CUBE_SIZE + GAP), + -z * (CUBE_SIZE + GAP) + ); + tensor.add(cube); + } + } + } + + console.log(`Highlighting ${coords.length} coordinates in global tensor`); + coords.forEach(([x, y, z]) => { + const cube = tensor.children.find(c => + c.userData.tensor0 === x && c.userData.tensor1 === y && c.userData.tensor2 === z + ); + if (cube) { + cube.material.color.set(COLOR_SLICE); + console.log(`Highlighted cube at (${x}, ${y}, ${z})`); + } else { + console.warn(`Could not find cube at (${x}, ${y}, ${z})`); + } + }); + } else { + console.log(`Creating slice tensor with ${coords.length} coordinates`); + coords.forEach(([x, y, z]) => { + const cube = createCube(color, tensorName, x, y, z, cubeGeometry, edgesGeometry, lineMaterial); + cube.position.set( + x * (CUBE_SIZE + GAP), + -y * (CUBE_SIZE + GAP), + -z * (CUBE_SIZE + GAP) + ); + tensor.add(cube); + }); + } + + console.log(`Created ${tensorName} tensor with ${tensor.children.length} cubes`); + return tensor; +} + +export function calculateTensorSize(shape) { + const [width, height, depth] = shape; + return new THREE.Vector3( + width * (CUBE_SIZE + GAP), + height * (CUBE_SIZE + GAP), + depth * (CUBE_SIZE + GAP) + ); +} + +export function interpolateColor(color1, color2, factor) { + return new THREE.Color().lerpColors(color1, color2, factor); +} + +export function updateCubeColor(tensor, coord, startColor, endColor, factor) { + const cube = tensor.children.find(c => + c.tensor0 === coord[0] && c.tensor1 === coord[1] && c.tensor2 === coord[2] + ); + if (cube) { + cube.material.color.copy(interpolateColor(startColor, endColor, factor)); + } +} + +export function setupCamera(scene, camera) { + const box = new THREE.Box3().setFromObject(scene); + const center = box.getCenter(new THREE.Vector3()); + const size = box.getSize(new THREE.Vector3()); + const maxDim = Math.max(size.x, size.y, size.z); + const fov = camera.fov * (Math.PI / 180); + let cameraZ = Math.abs(maxDim / 2 / Math.tan(fov / 2)); + cameraZ *= 1.5; + + camera.position.set(center.x, center.y, center.z + cameraZ); + camera.lookAt(center); + + return { center, cameraZ }; +} + +export function setupEventListeners(containerElement, camera, renderer, onMouseMove, onKeyDown) { + window.addEventListener('resize', () => { + camera.aspect = containerElement.clientWidth / containerElement.clientHeight; + camera.updateProjectionMatrix(); + renderer.setSize(containerElement.clientWidth, containerElement.clientHeight); + }); + containerElement.addEventListener('mousemove', onMouseMove); + window.addEventListener('keydown', onKeyDown); +} + +export function cameraControls(camera, cameraRotation) { + const PAN_SPEED = 0.1; + const TILT_SPEED = 0.02; + const ZOOM_SPEED = 0.5; + + return function onKeyDown(event) { + switch (event.key.toLowerCase()) { + case 'w': + camera.position.y += PAN_SPEED; + break; + case 's': + camera.position.y -= PAN_SPEED; + break; + case 'a': + camera.position.x -= PAN_SPEED; + break; + case 'd': + camera.position.x += PAN_SPEED; + break; + case 'arrowup': + cameraRotation.x -= TILT_SPEED; + break; + case 'arrowdown': + cameraRotation.x += TILT_SPEED; + break; + case 'arrowleft': + cameraRotation.y -= TILT_SPEED; + break; + case 'arrowright': + cameraRotation.y += TILT_SPEED; + break; + case 'o': + camera.position.z += ZOOM_SPEED; + break; + case 'p': + camera.position.z -= ZOOM_SPEED; + break; + case ' ': + break; + } + camera.setRotationFromEuler(cameraRotation); + camera.updateProjectionMatrix(); + }; +} diff --git a/triton_viz/static/matmul.js b/triton_viz/static/matmul.js new file mode 100644 index 0000000..d87d78b --- /dev/null +++ b/triton_viz/static/matmul.js @@ -0,0 +1,360 @@ +import * as THREE from 'https://cdn.jsdelivr.net/npm/three@0.155.0/build/three.module.js'; + +export function createMatMulVisualization(containerElement, op) { + const { input_shape, other_shape, output_shape } = op; + console.log(op.uuid) + fetch('/api/setop', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ uuid: op.uuid }), + }) + .then(response => response.json()) + .then(data => console.log('Set current op:', data)) + .catch((error) => console.error('Error:', error)); + let currentStep = 0; + const totalSteps = input_shape[1]; + let frame = 0; + + + const sideMenu = document.createElement('div'); + sideMenu.style.position = 'absolute'; + sideMenu.style.top = '10px'; + sideMenu.style.right = '10px'; + sideMenu.style.width = '200px'; + sideMenu.style.padding = '10px'; + sideMenu.style.backgroundColor = 'rgba(0, 0, 0, 0.7)'; + sideMenu.style.color = 'white'; + sideMenu.style.fontFamily = 'Arial, sans-serif'; + sideMenu.style.fontSize = '14px'; + sideMenu.style.borderRadius = '5px'; + containerElement.appendChild(sideMenu); + let hoveredCube = null; + + + + + async function getElementValue( matrixName, row, col) { + let uuid = op.uuid; + const response = await fetch('/api/getValue', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ uuid, matrixName, row, col, currentStep }), + }); + return await response.json(); + } + + + function updateSideMenu(matrix, x, y) { + if (!matrix) { + sideMenu.innerHTML = ''; + return; + } + + let matrixName; + let dims; + if (matrix === matrixA) { + matrixName = 'A'; + dims = input_shape; + } else if (matrix === matrixB) { + matrixName = 'B'; + dims = other_shape; + } else if (matrix === matrixC) { + matrixName = 'C'; + dims = output_shape; + } else { + sideMenu.innerHTML = ''; + return; + } + console.log(matrixName, "x:", (x + 1), "y:", (y + 1)); + sideMenu.innerHTML = ` +Row: ${y + 1}
+Column: ${x + 1}
+Dimensions: ${dims[0]} x ${dims[1]}
+ `; + } + + const raycaster = new THREE.Raycaster(); + const mouse = new THREE.Vector2(); + + + async function onMouseMove(event) { + mouse.x = (event.clientX / containerElement.clientWidth) * 2 - 1; + mouse.y = -(event.clientY / containerElement.clientHeight) * 2 + 1; + + raycaster.setFromCamera(mouse, camera); + + const allMatrixChildren = [ + ...(matrixA ? matrixA.children : []), + ...(matrixB ? matrixB.children : []), + ...(matrixC ? matrixC.children : []) + ]; + + const intersects = raycaster.intersectObjects(allMatrixChildren, true); + + if (hoveredCube) { + hoveredCube.getObjectByName('hoverOutline').visible = false; + hoveredCube = null; + } + + if (intersects.length > 0) { + // Find the actual cube (parent of the intersected object) + hoveredCube = intersects[0].object; + while (hoveredCube && !hoveredCube.matrixName) { + hoveredCube = hoveredCube.parent; + } + + if (hoveredCube) { + const hoverOutline = hoveredCube.getObjectByName('hoverOutline'); + if (hoverOutline) { + hoverOutline.visible = true; + } + const res = await getElementValue(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol); + // Log the matrix name, row, and column of the hovered cube + console.log( + // `Matrix: ${hoveredCube.matrixName}, ` + + // `Row: ${hoveredCube.matrixRow + 1}, ` + + // `Column: ${hoveredCube.matrixCol + 1}`+ + `Value: ${res.value}` + ); + + updateSideMenu(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol); + } + } else { + updateSideMenu(null); + } + } + + + + + + const CUBE_SIZE = 0.2; + const GAP = 0.05; + + const COLOR_A = new THREE.Color(0.53, 0.81, 0.98); + const COLOR_B = new THREE.Color(1.0, 0.65, 0.0); + const COLOR_C = new THREE.Color(1.0, 1.0, 1.0); + const COLOR_HIGHLIGHT = new THREE.Color(0.0, 0.0, 1.0); + const COLOR_FILLED = new THREE.Color(0.0, 0.0, 1.0); + const COLOR_BACKGROUND = new THREE.Color(0.0, 0.0, 0.0); + const COLOR_EDGE = new THREE.Color(0.3, 0.3, 0.3); + const COLOR_HOVER = new THREE.Color(1.0, 1.0, 0.0); + + const scene = new THREE.Scene(); + scene.background = COLOR_BACKGROUND; + const camera = new THREE.PerspectiveCamera(45, containerElement.clientWidth / containerElement.clientHeight, 0.1, 1000); + + const renderer = new THREE.WebGLRenderer({ antialias: true }); + renderer.setSize(containerElement.clientWidth, containerElement.clientHeight); + containerElement.appendChild(renderer.domElement); + + const ambientLight = new THREE.AmbientLight(0xffffff, 0.5); + scene.add(ambientLight); + const directionalLight = new THREE.DirectionalLight(0xffffff, 0.5); + directionalLight.position.set(10, 10, 10); + scene.add(directionalLight); + + const cubeGeometry = new THREE.BoxGeometry(CUBE_SIZE, CUBE_SIZE, CUBE_SIZE); + const edgesGeometry = new THREE.EdgesGeometry(cubeGeometry); + const lineMaterial = new THREE.LineBasicMaterial({ color: COLOR_EDGE }); + + function createCube(color, matrixName, i, j) { + const cubeMaterial = new THREE.MeshPhongMaterial({ color: color }); + const cube = new THREE.Mesh(cubeGeometry, cubeMaterial); + const edges = new THREE.LineSegments(edgesGeometry, lineMaterial); + cube.add(edges); + + const hoverGeometry = new THREE.BoxGeometry(CUBE_SIZE * 1.05, CUBE_SIZE * 1.05, CUBE_SIZE * 1.05); + const hoverEdgesGeometry = new THREE.EdgesGeometry(hoverGeometry); + const hoverOutline = new THREE.LineSegments(hoverEdgesGeometry, new THREE.LineBasicMaterial({ color: COLOR_HOVER })); + hoverOutline.visible = false; + hoverOutline.name = 'hoverOutline'; + cube.add(hoverOutline); + + cube.name = `${matrixName}_cube_${i}_${j}`; + cube.matrixName = matrixName; + cube.matrixRow = i; + cube.matrixCol = j; + + return cube; + } + + function createMatrix(dimensions, position, color, matrixName) { + const matrix = new THREE.Group(); + matrix.userData.dimensions = dimensions; + for (let i = 0; i < dimensions[0]; i++) { + for (let j = 0; j < dimensions[1]; j++) { + const cube = createCube(color, matrixName, i, j); + cube.position.set( + position.x + j * (CUBE_SIZE + GAP), + position.y - i * (CUBE_SIZE + GAP), + position.z + ); + matrix.add(cube); + } + } + return matrix; + } + + const matrixA = createMatrix(input_shape, new THREE.Vector3(-10, 10, 0), COLOR_A, 'A'); + const matrixB = createMatrix(other_shape, new THREE.Vector3(0, 10, 0), COLOR_B, 'B'); + const matrixC = createMatrix(output_shape, new THREE.Vector3(-5, -4, 0), COLOR_C, 'C'); + + scene.add(matrixA); + scene.add(matrixB); + scene.add(matrixC); + + const center = new THREE.Vector3(); + const size = new THREE.Vector3(); + const box = new THREE.Box3().setFromObject(scene); + box.getCenter(center); + box.getSize(size); + const maxDim = Math.max(size.x, size.y, size.z); + const fov = camera.fov * (Math.PI / 180); + let cameraZ = Math.abs(maxDim / 2 / Math.tan(fov / 2)); + cameraZ *= 1.5; + + camera.position.set(center.x, center.y, center.z + cameraZ); + camera.lookAt(center); + + let isPaused = false; + + const totalFrames = input_shape[0] * other_shape[1]; + + function highlightCubes(matrix, indices, highlightColor) { + indices.forEach(([i, j]) => { + if (i >= 0 && i < matrix.userData.dimensions[0] && j >= 0 && j < matrix.userData.dimensions[1]) { + const index = i * matrix.userData.dimensions[1] + j; + if (index < matrix.children.length) { + matrix.children[index].material.color.copy(highlightColor); + + } + } + }); + } + + function resetColors() { + matrixA.children.forEach(cube => cube.material.color.copy(COLOR_A)); + matrixB.children.forEach(cube => cube.material.color.copy(COLOR_B)); + } + + function animate() { + requestAnimationFrame(animate); + + if (!isPaused && frame < totalFrames) { + resetColors(); + + const row = Math.floor(frame / other_shape[1]); + const col = frame % other_shape[1]; + currentStep = frame % totalSteps + 1; + + const highlightA = Array.from({ length: input_shape[1] }, (_, i) => [row, i]); + const highlightB = Array.from({ length: other_shape[0] }, (_, i) => [i, col]); + const highlightC = [[row, col]]; + + highlightCubes(matrixA, highlightA, COLOR_HIGHLIGHT); + highlightCubes(matrixB, highlightB, COLOR_HIGHLIGHT); + highlightCubes(matrixC, highlightC, COLOR_FILLED); + + frame++; + } + + renderer.render(scene, camera); + } + + function onResize() { + camera.aspect = containerElement.clientWidth / containerElement.clientHeight; + camera.updateProjectionMatrix(); + renderer.setSize(containerElement.clientWidth, containerElement.clientHeight); + } + + const controlPanel = document.createElement('div'); + controlPanel.style.position = 'absolute'; + controlPanel.style.bottom = '10px'; + controlPanel.style.left = '10px'; + controlPanel.style.display = 'flex'; + controlPanel.style.gap = '10px'; + + const playPauseButton = document.createElement('button'); + playPauseButton.textContent = 'Play/Pause'; + playPauseButton.addEventListener('click', () => { + isPaused = !isPaused; + }); + + const resetButton = document.createElement('button'); + resetButton.textContent = 'Reset'; + resetButton.addEventListener('click', () => { + frame = 0; + resetColors(); + }); + + controlPanel.appendChild(playPauseButton); + controlPanel.appendChild(resetButton); + + containerElement.appendChild(controlPanel); + + const PAN_SPEED = 0.1; + const TILT_SPEED = 0.02; + const ZOOM_SPEED = 0.5; + let cameraRotation = new THREE.Euler(0, 0, 0, 'YXZ'); + + function onKeyDown(event) { + switch (event.key.toLowerCase()) { + case 'w': + camera.position.y += PAN_SPEED; + break; + case 's': + camera.position.y -= PAN_SPEED; + break; + case 'a': + camera.position.x -= PAN_SPEED; + break; + case 'd': + camera.position.x += PAN_SPEED; + break; + case 'arrowup': + cameraRotation.x -= TILT_SPEED; + break; + case 'arrowdown': + cameraRotation.x += TILT_SPEED; + break; + case 'arrowleft': + cameraRotation.y -= TILT_SPEED; + break; + case 'arrowright': + cameraRotation.y += TILT_SPEED; + break; + case 'o': + camera.position.z += ZOOM_SPEED; + break; + case 'p': + camera.position.z -= ZOOM_SPEED; + break; + } + camera.setRotationFromEuler(cameraRotation); + camera.updateProjectionMatrix(); + } + + window.addEventListener('resize', onResize); + window.addEventListener('keydown', onKeyDown); + containerElement.addEventListener('mousemove', onMouseMove); + + animate(); + + + function cleanup() { + window.removeEventListener('resize', onResize); + window.removeEventListener('keydown', onKeyDown); + containerElement.removeEventListener('mousemove', onMouseMove); + containerElement.innerHTML = ''; + renderer.dispose(); + scene.clear(); + } + + return cleanup; +} diff --git a/triton_viz/static/store-utils.js b/triton_viz/static/store-utils.js new file mode 100644 index 0000000..d87d78b --- /dev/null +++ b/triton_viz/static/store-utils.js @@ -0,0 +1,360 @@ +import * as THREE from 'https://cdn.jsdelivr.net/npm/three@0.155.0/build/three.module.js'; + +export function createMatMulVisualization(containerElement, op) { + const { input_shape, other_shape, output_shape } = op; + console.log(op.uuid) + fetch('/api/setop', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ uuid: op.uuid }), + }) + .then(response => response.json()) + .then(data => console.log('Set current op:', data)) + .catch((error) => console.error('Error:', error)); + let currentStep = 0; + const totalSteps = input_shape[1]; + let frame = 0; + + + const sideMenu = document.createElement('div'); + sideMenu.style.position = 'absolute'; + sideMenu.style.top = '10px'; + sideMenu.style.right = '10px'; + sideMenu.style.width = '200px'; + sideMenu.style.padding = '10px'; + sideMenu.style.backgroundColor = 'rgba(0, 0, 0, 0.7)'; + sideMenu.style.color = 'white'; + sideMenu.style.fontFamily = 'Arial, sans-serif'; + sideMenu.style.fontSize = '14px'; + sideMenu.style.borderRadius = '5px'; + containerElement.appendChild(sideMenu); + let hoveredCube = null; + + + + + async function getElementValue( matrixName, row, col) { + let uuid = op.uuid; + const response = await fetch('/api/getValue', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ uuid, matrixName, row, col, currentStep }), + }); + return await response.json(); + } + + + function updateSideMenu(matrix, x, y) { + if (!matrix) { + sideMenu.innerHTML = ''; + return; + } + + let matrixName; + let dims; + if (matrix === matrixA) { + matrixName = 'A'; + dims = input_shape; + } else if (matrix === matrixB) { + matrixName = 'B'; + dims = other_shape; + } else if (matrix === matrixC) { + matrixName = 'C'; + dims = output_shape; + } else { + sideMenu.innerHTML = ''; + return; + } + console.log(matrixName, "x:", (x + 1), "y:", (y + 1)); + sideMenu.innerHTML = ` +Row: ${y + 1}
+Column: ${x + 1}
+Dimensions: ${dims[0]} x ${dims[1]}
+ `; + } + + const raycaster = new THREE.Raycaster(); + const mouse = new THREE.Vector2(); + + + async function onMouseMove(event) { + mouse.x = (event.clientX / containerElement.clientWidth) * 2 - 1; + mouse.y = -(event.clientY / containerElement.clientHeight) * 2 + 1; + + raycaster.setFromCamera(mouse, camera); + + const allMatrixChildren = [ + ...(matrixA ? matrixA.children : []), + ...(matrixB ? matrixB.children : []), + ...(matrixC ? matrixC.children : []) + ]; + + const intersects = raycaster.intersectObjects(allMatrixChildren, true); + + if (hoveredCube) { + hoveredCube.getObjectByName('hoverOutline').visible = false; + hoveredCube = null; + } + + if (intersects.length > 0) { + // Find the actual cube (parent of the intersected object) + hoveredCube = intersects[0].object; + while (hoveredCube && !hoveredCube.matrixName) { + hoveredCube = hoveredCube.parent; + } + + if (hoveredCube) { + const hoverOutline = hoveredCube.getObjectByName('hoverOutline'); + if (hoverOutline) { + hoverOutline.visible = true; + } + const res = await getElementValue(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol); + // Log the matrix name, row, and column of the hovered cube + console.log( + // `Matrix: ${hoveredCube.matrixName}, ` + + // `Row: ${hoveredCube.matrixRow + 1}, ` + + // `Column: ${hoveredCube.matrixCol + 1}`+ + `Value: ${res.value}` + ); + + updateSideMenu(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol); + } + } else { + updateSideMenu(null); + } + } + + + + + + const CUBE_SIZE = 0.2; + const GAP = 0.05; + + const COLOR_A = new THREE.Color(0.53, 0.81, 0.98); + const COLOR_B = new THREE.Color(1.0, 0.65, 0.0); + const COLOR_C = new THREE.Color(1.0, 1.0, 1.0); + const COLOR_HIGHLIGHT = new THREE.Color(0.0, 0.0, 1.0); + const COLOR_FILLED = new THREE.Color(0.0, 0.0, 1.0); + const COLOR_BACKGROUND = new THREE.Color(0.0, 0.0, 0.0); + const COLOR_EDGE = new THREE.Color(0.3, 0.3, 0.3); + const COLOR_HOVER = new THREE.Color(1.0, 1.0, 0.0); + + const scene = new THREE.Scene(); + scene.background = COLOR_BACKGROUND; + const camera = new THREE.PerspectiveCamera(45, containerElement.clientWidth / containerElement.clientHeight, 0.1, 1000); + + const renderer = new THREE.WebGLRenderer({ antialias: true }); + renderer.setSize(containerElement.clientWidth, containerElement.clientHeight); + containerElement.appendChild(renderer.domElement); + + const ambientLight = new THREE.AmbientLight(0xffffff, 0.5); + scene.add(ambientLight); + const directionalLight = new THREE.DirectionalLight(0xffffff, 0.5); + directionalLight.position.set(10, 10, 10); + scene.add(directionalLight); + + const cubeGeometry = new THREE.BoxGeometry(CUBE_SIZE, CUBE_SIZE, CUBE_SIZE); + const edgesGeometry = new THREE.EdgesGeometry(cubeGeometry); + const lineMaterial = new THREE.LineBasicMaterial({ color: COLOR_EDGE }); + + function createCube(color, matrixName, i, j) { + const cubeMaterial = new THREE.MeshPhongMaterial({ color: color }); + const cube = new THREE.Mesh(cubeGeometry, cubeMaterial); + const edges = new THREE.LineSegments(edgesGeometry, lineMaterial); + cube.add(edges); + + const hoverGeometry = new THREE.BoxGeometry(CUBE_SIZE * 1.05, CUBE_SIZE * 1.05, CUBE_SIZE * 1.05); + const hoverEdgesGeometry = new THREE.EdgesGeometry(hoverGeometry); + const hoverOutline = new THREE.LineSegments(hoverEdgesGeometry, new THREE.LineBasicMaterial({ color: COLOR_HOVER })); + hoverOutline.visible = false; + hoverOutline.name = 'hoverOutline'; + cube.add(hoverOutline); + + cube.name = `${matrixName}_cube_${i}_${j}`; + cube.matrixName = matrixName; + cube.matrixRow = i; + cube.matrixCol = j; + + return cube; + } + + function createMatrix(dimensions, position, color, matrixName) { + const matrix = new THREE.Group(); + matrix.userData.dimensions = dimensions; + for (let i = 0; i < dimensions[0]; i++) { + for (let j = 0; j < dimensions[1]; j++) { + const cube = createCube(color, matrixName, i, j); + cube.position.set( + position.x + j * (CUBE_SIZE + GAP), + position.y - i * (CUBE_SIZE + GAP), + position.z + ); + matrix.add(cube); + } + } + return matrix; + } + + const matrixA = createMatrix(input_shape, new THREE.Vector3(-10, 10, 0), COLOR_A, 'A'); + const matrixB = createMatrix(other_shape, new THREE.Vector3(0, 10, 0), COLOR_B, 'B'); + const matrixC = createMatrix(output_shape, new THREE.Vector3(-5, -4, 0), COLOR_C, 'C'); + + scene.add(matrixA); + scene.add(matrixB); + scene.add(matrixC); + + const center = new THREE.Vector3(); + const size = new THREE.Vector3(); + const box = new THREE.Box3().setFromObject(scene); + box.getCenter(center); + box.getSize(size); + const maxDim = Math.max(size.x, size.y, size.z); + const fov = camera.fov * (Math.PI / 180); + let cameraZ = Math.abs(maxDim / 2 / Math.tan(fov / 2)); + cameraZ *= 1.5; + + camera.position.set(center.x, center.y, center.z + cameraZ); + camera.lookAt(center); + + let isPaused = false; + + const totalFrames = input_shape[0] * other_shape[1]; + + function highlightCubes(matrix, indices, highlightColor) { + indices.forEach(([i, j]) => { + if (i >= 0 && i < matrix.userData.dimensions[0] && j >= 0 && j < matrix.userData.dimensions[1]) { + const index = i * matrix.userData.dimensions[1] + j; + if (index < matrix.children.length) { + matrix.children[index].material.color.copy(highlightColor); + + } + } + }); + } + + function resetColors() { + matrixA.children.forEach(cube => cube.material.color.copy(COLOR_A)); + matrixB.children.forEach(cube => cube.material.color.copy(COLOR_B)); + } + + function animate() { + requestAnimationFrame(animate); + + if (!isPaused && frame < totalFrames) { + resetColors(); + + const row = Math.floor(frame / other_shape[1]); + const col = frame % other_shape[1]; + currentStep = frame % totalSteps + 1; + + const highlightA = Array.from({ length: input_shape[1] }, (_, i) => [row, i]); + const highlightB = Array.from({ length: other_shape[0] }, (_, i) => [i, col]); + const highlightC = [[row, col]]; + + highlightCubes(matrixA, highlightA, COLOR_HIGHLIGHT); + highlightCubes(matrixB, highlightB, COLOR_HIGHLIGHT); + highlightCubes(matrixC, highlightC, COLOR_FILLED); + + frame++; + } + + renderer.render(scene, camera); + } + + function onResize() { + camera.aspect = containerElement.clientWidth / containerElement.clientHeight; + camera.updateProjectionMatrix(); + renderer.setSize(containerElement.clientWidth, containerElement.clientHeight); + } + + const controlPanel = document.createElement('div'); + controlPanel.style.position = 'absolute'; + controlPanel.style.bottom = '10px'; + controlPanel.style.left = '10px'; + controlPanel.style.display = 'flex'; + controlPanel.style.gap = '10px'; + + const playPauseButton = document.createElement('button'); + playPauseButton.textContent = 'Play/Pause'; + playPauseButton.addEventListener('click', () => { + isPaused = !isPaused; + }); + + const resetButton = document.createElement('button'); + resetButton.textContent = 'Reset'; + resetButton.addEventListener('click', () => { + frame = 0; + resetColors(); + }); + + controlPanel.appendChild(playPauseButton); + controlPanel.appendChild(resetButton); + + containerElement.appendChild(controlPanel); + + const PAN_SPEED = 0.1; + const TILT_SPEED = 0.02; + const ZOOM_SPEED = 0.5; + let cameraRotation = new THREE.Euler(0, 0, 0, 'YXZ'); + + function onKeyDown(event) { + switch (event.key.toLowerCase()) { + case 'w': + camera.position.y += PAN_SPEED; + break; + case 's': + camera.position.y -= PAN_SPEED; + break; + case 'a': + camera.position.x -= PAN_SPEED; + break; + case 'd': + camera.position.x += PAN_SPEED; + break; + case 'arrowup': + cameraRotation.x -= TILT_SPEED; + break; + case 'arrowdown': + cameraRotation.x += TILT_SPEED; + break; + case 'arrowleft': + cameraRotation.y -= TILT_SPEED; + break; + case 'arrowright': + cameraRotation.y += TILT_SPEED; + break; + case 'o': + camera.position.z += ZOOM_SPEED; + break; + case 'p': + camera.position.z -= ZOOM_SPEED; + break; + } + camera.setRotationFromEuler(cameraRotation); + camera.updateProjectionMatrix(); + } + + window.addEventListener('resize', onResize); + window.addEventListener('keydown', onKeyDown); + containerElement.addEventListener('mousemove', onMouseMove); + + animate(); + + + function cleanup() { + window.removeEventListener('resize', onResize); + window.removeEventListener('keydown', onKeyDown); + containerElement.removeEventListener('mousemove', onMouseMove); + containerElement.innerHTML = ''; + renderer.dispose(); + scene.clear(); + } + + return cleanup; +} diff --git a/triton_viz/static/store.js b/triton_viz/static/store.js new file mode 100644 index 0000000..470b327 --- /dev/null +++ b/triton_viz/static/store.js @@ -0,0 +1,221 @@ +import * as THREE from 'https://cdn.jsdelivr.net/npm/three@0.155.0/build/three.module.js'; +import { + setupScene, + setupGeometries, + createTensor, + calculateTensorSize, + updateCubeColor, + setupCamera, + setupEventListeners, + cameraControls +} from './store-utils.js'; + +export function createStoreVisualization(containerElement, op) { + + console.log(op.uuid); + fetch('/api/setop', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ uuid: op.uuid }), + }) + .then(response => response.json()) + .then(data => console.log('Set current op:', data)) + .catch((error) => console.error('Error:', error)); + + let currentStep = 0; + let frame = 0; + let isPaused = false; + + const sideMenu = createSideMenu(containerElement); + let hoveredCube = null; + + const COLOR_GLOBAL = new THREE.Color(0.2, 0.2, 0.2); // Dark Gray + const COLOR_SLICE = new THREE.Color(0.0, 0.7, 1.0); // Cyan (starting color for global slice) + const COLOR_LEFT_SLICE = new THREE.Color(1.0, 0.0, 1.0); // Magenta (starting color for left slice) + const COLOR_LOADED = new THREE.Color(1.0, 0.8, 0.0); // Gold (final color for both slices) + const COLOR_BACKGROUND = new THREE.Color(0.0, 0.0, 0.0); // Black + + const { scene, camera, renderer } = setupScene(containerElement, COLOR_BACKGROUND); + const { cubeGeometry, edgesGeometry, lineMaterial } = setupGeometries(); + + const globalTensor = createTensor(op.global_shape, op.global_coords, COLOR_GLOBAL, 'Global', cubeGeometry, edgesGeometry, lineMaterial); + const sliceTensor = createTensor(op.slice_shape, op.slice_coords, COLOR_LEFT_SLICE, 'Slice', cubeGeometry, edgesGeometry, lineMaterial); + + // Position slice tensor + const globalSize = calculateTensorSize(op.global_shape); + sliceTensor.position.set(globalSize.x + 5, 0, 0); // Adjusted tensor spacing + + scene.add(globalTensor); + scene.add(sliceTensor); + + addLabels(scene, globalTensor, sliceTensor); + setupCamera(scene, camera); + + const totalFrames = op.global_coords.length * 2 + 30; + + const raycaster = new THREE.Raycaster(); + const mouse = new THREE.Vector2(); + + const onKeyDown = cameraControls(camera, new THREE.Euler(0, 0, 0, 'YXZ')); + setupEventListeners(containerElement, camera, renderer, onMouseMove, onKeyDown); + animate(); + + async function onMouseMove(event) { + mouse.x = (event.clientX / containerElement.clientWidth) * 2 - 1; + mouse.y = -(event.clientY / containerElement.clientHeight) * 2 + 1; + + raycaster.setFromCamera(mouse, camera); + + const allTensorChildren = [ + ...globalTensor.children, + ...sliceTensor.children + ]; + + const intersects = raycaster.intersectObjects(allTensorChildren, true); + + if (hoveredCube) { + hoveredCube.getObjectByName('hoverOutline').visible = false; + hoveredCube = null; + } + + if (intersects.length > 0) { + hoveredCube = intersects[0].object; + while (hoveredCube && !hoveredCube.tensorName) { + hoveredCube = hoveredCube.parent; + } + + if (hoveredCube) { + const hoverOutline = hoveredCube.getObjectByName('hoverOutline'); + if (hoverOutline) { + hoverOutline.visible = true; + } + + updateSideMenu(hoveredCube.tensorName, hoveredCube.tensor0, hoveredCube.tensor1, hoveredCube.tensor2, undefined); + + const res = await getElementValue(hoveredCube.tensorName, hoveredCube.tensor0, hoveredCube.tensor1, hoveredCube.tensor2); + + updateSideMenu(hoveredCube.tensorName, hoveredCube.tensor0, hoveredCube.tensor1, hoveredCube.tensor2, res.value); + + console.log(`Value: ${res.value}`); + } + } else { + updateSideMenu(null); + } + } + + function animate() { + requestAnimationFrame(animate); + + if (!isPaused && frame < totalFrames) { + const index = Math.floor(frame / 2); + const factor = (frame % 2) / 1.0; + + if (index < op.global_coords.length) { + const globalCoord = op.global_coords[index]; + const sliceCoord = op.slice_coords[index]; + + updateCubeColor(globalTensor, globalCoord, COLOR_GLOBAL, COLOR_SLICE, factor); + updateCubeColor(sliceTensor, sliceCoord, COLOR_LEFT_SLICE, COLOR_LOADED, factor); + + highlightCurrentOperation(globalTensor, globalCoord, sliceTensor, sliceCoord); + updateInfoPanel(globalCoord, sliceCoord, index); + } + + frame++; + } + + renderer.render(scene, camera); + } + + function highlightCurrentOperation(globalTensor, globalCoord, sliceTensor, sliceCoord) { + globalTensor.children.forEach(cube => cube.material.emissive.setHex(0x000000)); + sliceTensor.children.forEach(cube => cube.material.emissive.setHex(0x000000)); + + const globalCube = globalTensor.children.find(c => + c.tensor0 === globalCoord[0] && c.tensor1 === globalCoord[1] && c.tensor2 === globalCoord[2] + ); + const sliceCube = sliceTensor.children.find(c => + c.tensor0 === sliceCoord[0] && c.tensor1 === sliceCoord[1] && c.tensor2 === sliceCoord[2] + ); + + if (globalCube) globalCube.material.emissive.setHex(0x444444); + if (sliceCube) sliceCube.material.emissive.setHex(0x444444); + } + + async function getElementValue(tensorName, x, y, z) { + let uuid = op.uuid; + const response = await fetch('/api/getStoreValue', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ uuid, tensorName, x, y, z }), + }); + return await response.json(); + } + + function updateSideMenu(tensorName, x, y, z, value) { + if (!tensorName) { + sideMenu.innerHTML = ''; + return; + } + + let dims = tensorName === 'Global' ? op.global_shape : op.slice_shape; + sideMenu.innerHTML = ` +X: ${x + 1}
+Y: ${y + 1}
+Z: ${z + 1}
+Dimensions: ${dims.join(' x ')}
+Value: ${value !== undefined ? value : 'Storeing...'}
+ `; + } + + function updateInfoPanel(globalCoord, sliceCoord, index) { + sideMenu.innerHTML = ` +Global Coords: (${globalCoord.join(', ')})
+Slice Coords: (${sliceCoord.join(', ')})
+Progress: ${index + 1}/${op.global_coords.length}
+ `; + } + + function createSideMenu(container) { + const menu = document.createElement('div'); + menu.style.position = 'absolute'; + menu.style.top = '10px'; + menu.style.right = '10px'; + menu.style.width = '200px'; + menu.style.padding = '10px'; + menu.style.backgroundColor = 'rgba(0, 0, 0, 0.7)'; + menu.style.color = 'white'; + menu.style.fontFamily = 'Arial, sans-serif'; + menu.style.fontSize = '14px'; + menu.style.borderRadius = '5px'; + container.appendChild(menu); + return menu; + } + + function addLabels(scene, globalTensor, sliceTensor) { + addLabel(scene, "Global Tensor", globalTensor.position); + addLabel(scene, "Slice Tensor", sliceTensor.position); + } + + function addLabel(scene, text, position) { + const canvas = document.createElement('canvas'); + const context = canvas.getContext('2d'); + context.font = 'Bold 24px Arial'; + context.fillStyle = 'white'; + context.fillText(text, 0, 24); + + const texture = new THREE.CanvasTexture(canvas); + const material = new THREE.SpriteMaterial({ map: texture }); + const sprite = new THREE.Sprite(material); + sprite.position.set(position.x, position.y + 2, position.z); + sprite.scale.set(4, 2, 1); + scene.add(sprite); + } + +} diff --git a/triton_viz/static/visualization.js b/triton_viz/static/visualization.js new file mode 100644 index 0000000..e0d157c --- /dev/null +++ b/triton_viz/static/visualization.js @@ -0,0 +1,389 @@ +import { GridBlock } from './gridblock.js'; +import { createInfoPopup, showInfoPopup } from './infoPopup.js'; +let globalData; +let currentView = 'main'; +let canvas, ctx; +let maxX = 0, maxY = 0, maxZ = 0; +let sliders = [], zSlider, precomputeButton, kernelGrid; +let backButton; +let currentBlockData = null; +let isInitialized = false; +let containerElement; +let infoPopup; +let infoButton; +function switchToMainView() { + currentView = 'main'; + if (currentBlockData) { + currentBlockData.hideDetailedView(); + currentBlockData = null; + } + containerElement.style.pointerEvents = 'none'; + containerElement.style.display = 'none'; + containerElement.innerHTML = ''; + + canvas.style.display = 'block'; + draw(); +} + +function initializeApp() { + canvas = document.getElementById('canvas'); + if (!canvas) { + console.error('Canvas element not found'); + return; + } + ctx = canvas.getContext('2d'); + if (!ctx) { + console.error('Unable to get 2D context from canvas'); + return; + } + canvas.width = window.innerWidth; + canvas.height = window.innerHeight; + + canvas.addEventListener('mousedown', handleMouseEvent); + canvas.addEventListener('mouseup', handleMouseEvent); + canvas.addEventListener('mousemove', handleMouseEvent); + + containerElement = document.getElementById('visualization-container'); + if (!containerElement) { + console.error('Visualization container element not found'); + return; + } + + containerElement.style.pointerEvents = 'none'; + containerElement.style.display = 'none'; + + fetchData(); +} + +class Slider { + constructor(x, y, width, height, label, min_value = -1, max_value = 100) { + this.rect = { x, y, width, height }; + this.label = label; + this.min = min_value; + this.max = max_value; + this.value = min_value; + this.grabbed = false; + this.enabled = true; + } + + draw(ctx) { + if (!this.enabled) return; + ctx.fillStyle = '#3c3c46'; + ctx.fillRect(this.rect.x, this.rect.y, this.rect.width, this.rect.height); + const buttonX = this.rect.x + (this.value - this.min) / (this.max - this.min) * this.rect.width; + ctx.fillStyle = '#c8c8c8'; + ctx.fillRect(buttonX - 5, this.rect.y - 2, 10, this.rect.height + 4); + + ctx.fillStyle = '#c8c8c8'; + ctx.font = '18px Arial'; + ctx.fillText(this.label, this.rect.x, this.rect.y - 10); + ctx.fillText(this.value.toString(), this.rect.x + this.rect.width + 10, this.rect.y + this.rect.height / 2 + 5); + } + + handleEvent(event) { + if (!this.enabled) return; + if (event.type === 'mousedown') { + if (this.isPointInside(event.offsetX, event.offsetY)) { + this.grabbed = true; + } + } else if (event.type === 'mouseup') { + this.grabbed = false; + } else if (event.type === 'mousemove' && this.grabbed) { + const mouseX = event.offsetX; + this.value = Math.round((mouseX - this.rect.x) / this.rect.width * (this.max - this.min) + this.min); + this.value = Math.max(this.min, Math.min(this.max, this.value)); + } + } + + isPointInside(x, y) { + return x >= this.rect.x && x <= this.rect.x + this.rect.width && + y >= this.rect.y && y <= this.rect.y + this.rect.height; + } +} + +class Button { + constructor(x, y, width, height, text, isIcon = false) { + this.rect = { x, y, width, height }; + this.text = text; + this.isIcon = isIcon; + this.color = '#3c3c46'; + this.hoverColor = '#50505a'; + this.clickColor = '#64646e'; + this.isHovered = false; + this.isClicked = false; + this.clickTime = 0; + } + + draw(ctx) { + let color = this.color; + if (this.isClicked && Date.now() - this.clickTime < 100) { + color = this.clickColor; + } else if (this.isHovered) { + color = this.hoverColor; + } + + ctx.fillStyle = color; + ctx.fillRect(this.rect.x, this.rect.y, this.rect.width, this.rect.height); + ctx.strokeStyle = '#c8c8c8'; + ctx.strokeRect(this.rect.x, this.rect.y, this.rect.width, this.rect.height); + + ctx.fillStyle = '#c8c8c8'; + ctx.font = this.isIcon ? 'bold 24px Arial' : '18px Arial'; + ctx.textAlign = 'center'; + ctx.textBaseline = 'middle'; + ctx.fillText(this.text, this.rect.x + this.rect.width / 2, this.rect.y + this.rect.height / 2); + } + + handleEvent(event) { + const { offsetX, offsetY } = event; + this.isHovered = this.isPointInside(offsetX, offsetY); + if (event.type === 'mousedown' && this.isHovered) { + this.isClicked = true; + this.clickTime = Date.now(); + console.log(`Button '${this.text}' clicked!`); + } else if (event.type === 'mouseup') { + this.isClicked = false; + } + } + + isPointInside(x, y) { + return x >= this.rect.x && x <= this.rect.x + this.rect.width && + y >= this.rect.y && y <= this.rect.y + this.rect.height; + } +} + +class KernelGrid { + constructor(x, y, width, height, gridSize, visualizationData) { + this.rect = { x, y, width, height }; + this.gridSize = gridSize; + this.visualizationData = visualizationData; + this.currentZ = 0; + this.blocks = []; + this.calculateBlockSize(); + this.createBlocks(); + this.selectedBlock = null; + this.filterValues = [-1, -1, -1]; // Default filter values for x, y, z + } + + calculateBlockSize() { + this.blockWidth = Math.floor(this.rect.width / this.gridSize[0]) - 1; + this.blockHeight = Math.floor(this.rect.height / this.gridSize[1]) - 1; + } + + createBlocks() { + this.blocks = []; + for (let y = 0; y < this.gridSize[1]; y++) { + for (let x = 0; x < this.gridSize[0]; x++) { + const blockX = this.rect.x + x * (this.blockWidth + 1); + const blockY = this.rect.y + y * (this.blockHeight + 1); + const gridKey = `(${x}, ${y}, ${this.currentZ})`; + const blockData = this.visualizationData[gridKey] || []; + const block = new GridBlock( + blockX, blockY, this.blockWidth, this.blockHeight, + x, y, this.currentZ, blockData, + switchToMainView, + containerElement, + canvas, + draw + ); + this.blocks.push(block); + } + } + } + + draw(ctx) { + ctx.fillStyle = '#F0F0F0'; + ctx.fillRect(this.rect.x, this.rect.y, this.rect.width, this.rect.height); + this.blocks.forEach(block => { + if (this.shouldDrawBlock(block)) { + block.draw(ctx); + } + }); + } + + shouldDrawBlock(block) { + return (this.filterValues[0] === -1 || block.gridPosition.x === this.filterValues[0]) && + (this.filterValues[1] === -1 || block.gridPosition.y === this.filterValues[1]) && + (this.filterValues[2] === -1 || block.gridPosition.z === this.filterValues[2]); + } + + updateZ(z) { + this.currentZ = z; + this.filterValues[2] = z; + this.blocks.forEach(block => { + block.gridPosition.z = z; + const gridKey = `(${block.gridPosition.x}, ${block.gridPosition.y}, ${z})`; + block.blockData = this.visualizationData[gridKey] || []; + }); + } + + handleClick(x, y) { + const clickedBlock = this.blocks.find(block => + block.isPointInside(x, y) && this.shouldDrawBlock(block) + ); + if (clickedBlock) { + console.log(`Clicked block at (${clickedBlock.gridPosition.x}, ${clickedBlock.gridPosition.y}, ${clickedBlock.gridPosition.z})`); + if (this.selectedBlock) { + this.selectedBlock.hideDetailedView(); + } + this.selectedBlock = clickedBlock; + clickedBlock.showDetailedView(); + return clickedBlock; + } + return null; + } + + handleMouseMove(x, y) { + this.blocks.forEach(block => { + if (this.shouldDrawBlock(block)) { + block.handleMouseMove(x, y); + } else { + block.isHovered = false; + } + }); + } + + updateFilter(dimension, value) { + this.filterValues[dimension] = value; + } +} + +function determineMaxValues(visualizationData) { + maxX = 0; + maxY = 0; + maxZ = 0; + const keys = Object.keys(visualizationData); + keys.forEach(key => { + const [x, y, z] = key.replace(/[()]/g, '').split(', ').map(Number); + if (x > maxX) maxX = x; + if (y > maxY) maxY = y; + if (z > maxZ) maxZ = z; + }); +} + +function initializeUIElements() { + sliders = [ + new Slider(1300, 50, 250, 20, "Program Id 0", -1, maxX), + new Slider(1300, 120, 250, 20, "Program Id 1", -1, maxY), + new Slider(1300, 190, 250, 20, "Program Id 2", -1, maxZ) + ]; + + zSlider = new Slider(50, 860, 1200, 20, "Z-axis", 0, maxZ); + zSlider.enabled = maxZ > 0; + + precomputeButton = new Button(1300, 260, 250, 40, "Precompute"); + kernelGrid = new KernelGrid(50, 50, 1200, 800, [maxX + 1, maxY + 1, maxZ + 1], globalData.ops.visualization_data); + backButton = new Button(50, 50, 100, 40, "Back"); + const buttonSize = 40; + const margin = 10; + infoButton = new Button( + canvas.width - buttonSize - margin, + margin, + buttonSize, + buttonSize, + "i", + true + ); + + isInitialized = true; + + infoPopup = createInfoPopup(); +} + +function switchToTensorView(clickedBlock) { + currentView = 'tensor'; + currentBlockData = clickedBlock; + console.log("Switched to tensor view. Block data:", clickedBlock); + + containerElement.style.pointerEvents = 'auto'; + containerElement.style.display = 'block'; + clickedBlock.showDetailedView(); + + canvas.style.display = 'none'; +} + +function handleMouseEvent(event) { + if (!isInitialized) { + console.warn('UI elements not initialized yet'); + return; + } + if (infoButton) { + infoButton.handleEvent(event); + if (event.type === 'mousedown' && infoButton.isHovered) { + showInfoPopup(infoPopup); + } + } + const { offsetX, offsetY } = event; + if (currentView === 'main') { + sliders.forEach((slider, index) => { + slider.handleEvent(event); + if (kernelGrid) { + kernelGrid.updateFilter(index, slider.value); + } + }); + if (zSlider && zSlider.enabled) { + zSlider.handleEvent(event); + if (kernelGrid) { + kernelGrid.updateZ(zSlider.value); + } + } + if (precomputeButton) { + precomputeButton.handleEvent(event); + } + if (kernelGrid) { + kernelGrid.handleMouseMove(offsetX, offsetY); + if (event.type === 'mousedown') { + const clickedBlock = kernelGrid.handleClick(offsetX, offsetY); + if (clickedBlock) { + switchToTensorView(clickedBlock); + } + } + } + } else if (currentView === 'tensor') { + if (backButton) { + backButton.handleEvent(event); + if (event.type === 'mousedown' && backButton.isHovered) { + switchToMainView(); + } + } + } + draw(); +} + +function draw() { + if (!ctx) { + console.error('Canvas context not available'); + return; + } + + ctx.fillStyle = '#1e1e28'; + ctx.fillRect(0, 0, canvas.width, canvas.height); + + if (currentView === 'main' || currentView === 'main') { + if (kernelGrid) kernelGrid.draw(ctx); + sliders.forEach(slider => slider.draw(ctx)); + if (zSlider && zSlider.enabled) { + zSlider.draw(ctx); + } + if (precomputeButton) precomputeButton.draw(ctx); + if (infoButton) { + infoButton.draw(ctx); + } + } +} + +async function fetchData() { + try { + const response = await fetch('/api/data'); + globalData = await response.json(); + console.log(globalData); + + determineMaxValues(globalData.ops.visualization_data); + initializeUIElements(); + draw(); + } catch (error) { + console.error('Error fetching data:', error); + } +} + +initializeApp(); \ No newline at end of file diff --git a/triton_viz/templates/index.html b/triton_viz/templates/index.html new file mode 100644 index 0000000..b8eaebe --- /dev/null +++ b/triton_viz/templates/index.html @@ -0,0 +1,47 @@ + + + + + +