Feature/device selector op details (#251)

tenstorrent · Nov 27, 2024 · f39727d · f39727d
2 parents 510ad00 + 3fce24a
commit f39727d
Show file tree

Hide file tree

Showing 10 changed files with 67 additions and 36 deletions.
diff --git a/src/components/DeviceSelector.tsx b/src/components/DeviceSelector.tsx
@@ -19,10 +19,10 @@ function DeviceSelector() {
     const [selectedDevice, setSelectedDevice] = useAtom(selectedDeviceAtom);
 
     const updateSelectedDevices = (device?: number) => {
-        setSelectedDevice(device !== selectedDevice ? device : undefined);
+        setSelectedDevice(device !== undefined && device !== selectedDevice ? device : null);
     };
 
-    return (
+    return devices && devices?.length > 1 ? (
         <Select
             className={classNames('device-selector', {
                 'has-selection': selectedDevice !== undefined,
@@ -46,10 +46,10 @@ function DeviceSelector() {
                 outlined
             />
         </Select>
-    );
+    ) : null;
 }
 
-const DeviceItem = (device: number, onClick: (device: number) => void, selectedDevice?: number) => {
+const DeviceItem = (device: number, onClick: (device: number) => void, selectedDevice: number | null) => {
     return (
         <li key={device}>
             <Checkbox

diff --git a/src/components/buffer-summary/BufferSummaryPlotRenderer.tsx b/src/components/buffer-summary/BufferSummaryPlotRenderer.tsx
@@ -88,21 +88,23 @@ function BufferSummaryPlotRenderer({ buffersByOperation, tensorListByOperation }
 
     return buffersByOperation && !isLoadingDevices && tensorListByOperation ? (
         <div className='buffer-summary-chart'>
-            <Switch
-                label='Buffer zoom'
-                checked={isZoomedIn}
-                onChange={() => {
-                    setIsZoomedIn(!isZoomedIn);
-                }}
-            />
-
-            <GlobalSwitch
-                label='Hex axis labels'
-                checked={showHex}
-                onChange={() => {
-                    setShowHex(!showHex);
-                }}
-            />
+            <div className='controls'>
+                <Switch
+                    label='Buffer zoom'
+                    checked={isZoomedIn}
+                    onChange={() => {
+                        setIsZoomedIn(!isZoomedIn);
+                    }}
+                />
+
+                <GlobalSwitch
+                    label='Hex axis labels'
+                    checked={showHex}
+                    onChange={() => {
+                        setShowHex(!showHex);
+                    }}
+                />
+            </div>
 
             <p className='x-axis-label'>Memory Address</p>
 

diff --git a/src/components/operation-details/OperationDetailsComponent.tsx b/src/components/operation-details/OperationDetailsComponent.tsx
@@ -5,14 +5,14 @@
 import React, { useState } from 'react';
 import { Button, ButtonGroup, Intent, Position, Switch, Tooltip } from '@blueprintjs/core';
 import { IconNames } from '@blueprintjs/icons';
-import { useAtom } from 'jotai';
+import { useAtom, useAtomValue } from 'jotai';
 import { useOperationDetails, useOperationsList, usePreviousOperationDetails } from '../../hooks/useAPI';
 import 'styles/components/OperationDetailsComponent.scss';
 import StackTrace from './StackTrace';
 import OperationDetailsNavigation from '../OperationDetailsNavigation';
 import { OperationDetails } from '../../model/OperationDetails';
 import { CONDENSED_PLOT_CHUNK_COLOR, PlotMouseEventCustom } from '../../definitions/PlotConfigurations';
-import { selectedAddressAtom, selectedTensorAtom, showHexAtom } from '../../store/app';
+import { selectedAddressAtom, selectedDeviceAtom, selectedTensorAtom, showHexAtom } from '../../store/app';
 import ProducerConsumersData from './ProducerConsumersData';
 import isValidNumber from '../../functions/isValidNumber';
 import TensorVisualisationComponent from '../tensor-sharding-visualization/TensorVisualisationComponent';
@@ -25,19 +25,21 @@ import TensorDetailsList from './TensorDetailsList';
 import OperationArguments from '../OperationArguments';
 import DeviceOperationsFullRender from './DeviceOperationsFullRender';
 import useBufferFocus from '../../hooks/useBufferFocus';
+import DeviceSelector from '../DeviceSelector';
 
 interface OperationDetailsProps {
     operationId: number;
 }
 
 const OperationDetailsComponent: React.FC<OperationDetailsProps> = ({ operationId }) => {
+    const selectedDevice = useAtomValue(selectedDeviceAtom);
     const { data: operations } = useOperationsList();
     const [zoomedInViewMainMemory, setZoomedInViewMainMemory] = useState(false);
     const [showCircularBuffer, setShowCircularBuffer] = useState(false);
     const [showHex, setShowHex] = useAtom(showHexAtom);
     const {
         operationDetails: { data: operationDetails, isLoading, status },
-    } = useOperationDetails(operationId);
+    } = useOperationDetails(operationId, selectedDevice);
 
     const { data: previousOperationDetails, isLoading: isPrevLoading } =
         usePreviousOperationDetails(operationId).operationDetails;
@@ -171,8 +173,11 @@ const OperationDetailsComponent: React.FC<OperationDetailsProps> = ({ operationI
                                     DRAM
                                 </Button>
                             </ButtonGroup>
+
+                            <DeviceSelector />
                         </div>
-                        <div className='zoom-controls'>
+
+                        <div className='controls'>
                             <Switch
                                 label='Buffer zoom'
                                 checked={zoomedInViewMainMemory}

diff --git a/src/components/report-selection/LocalFolderSelector.tsx b/src/components/report-selection/LocalFolderSelector.tsx
@@ -9,10 +9,10 @@ import { ChangeEvent, type FC, useEffect, useState } from 'react';
 import 'styles/components/FolderPicker.scss';
 import { useNavigate } from 'react-router';
 import { useQueryClient } from 'react-query';
-import { useAtom } from 'jotai';
+import { useAtom, useSetAtom } from 'jotai';
 import ROUTES from '../../definitions/routes';
 import useLocalConnection from '../../hooks/useLocal';
-import { reportLocationAtom } from '../../store/app';
+import { reportLocationAtom, selectedDeviceAtom } from '../../store/app';
 import { ConnectionStatus, ConnectionTestStates } from '../../definitions/ConnectionStatus';
 import FileStatusWrapper from '../FileStatusOverlayWrapper';
 import FileStatusOverlay from '../FileStatusOverlay';
@@ -36,6 +36,7 @@ const LocalFolderOptions: FC = () => {
     const navigate = useNavigate();
     const queryClient = useQueryClient();
     const [reportLocation, setReportLocation] = useAtom(reportLocationAtom);
+    const setSelectedDevice = useSetAtom(selectedDeviceAtom);
 
     const { uploadLocalFolder, checkRequiredFiles, filterReportFiles } = useLocalConnection();
     const [folderStatus, setFolderStatus] = useState<ConnectionStatus | undefined>();
@@ -98,6 +99,7 @@ const LocalFolderOptions: FC = () => {
     const viewOperation = () => {
         // keeping this here temporarily until proven otherwise
         queryClient.clear();
+        setSelectedDevice(0);
 
         navigate(ROUTES.OPERATIONS);
     };
@@ -149,13 +151,15 @@ const LocalFolderOptions: FC = () => {
                         />
                         <span className='bp5-file-upload-input'>{localUploadLabel}</span>
                     </label>
+
                     <Button
                         disabled={!isLocalReportMounted}
                         onClick={viewOperation}
                         icon={IconNames.EYE_OPEN}
                     >
                         View report
                     </Button>
+
                     <FileStatusWrapper>
                         {(fileProgress) => (
                             <FileStatusOverlay

diff --git a/src/components/report-selection/RemoteSyncConfigurator.tsx b/src/components/report-selection/RemoteSyncConfigurator.tsx
@@ -7,14 +7,14 @@ import { FC, useEffect, useState } from 'react';
 import { AnchorButton, Button, FormGroup, Tooltip } from '@blueprintjs/core';
 
 import { IconNames } from '@blueprintjs/icons';
-import { useAtom } from 'jotai';
+import { useAtom, useSetAtom } from 'jotai';
 import { useQueryClient } from 'react-query';
 import { useNavigate } from 'react-router';
 import { RemoteConnection, RemoteFolder } from '../../definitions/RemoteConnection';
 import ROUTES from '../../definitions/routes';
 import isRemoteFolderOutdated from '../../functions/isRemoteFolderOutdated';
 import useRemote from '../../hooks/useRemote';
-import { reportLocationAtom } from '../../store/app';
+import { reportLocationAtom, selectedDeviceAtom } from '../../store/app';
 import AddRemoteConnection from './AddRemoteConnection';
 import RemoteConnectionSelector from './RemoteConnectionSelector';
 import RemoteFolderSelector from './RemoteFolderSelector';
@@ -27,6 +27,7 @@ const RemoteSyncConfigurator: FC = () => {
         remote.persistentState.getSavedRemoteFolders(remote.persistentState.selectedConnection),
     );
     const [_, setReportLocation] = useAtom(reportLocationAtom);
+    const setSelectedDevice = useSetAtom(selectedDeviceAtom);
 
     const [isSyncingRemoteFolder, setIsSyncingRemoteFolder] = useState(false);
     const [isLoadingFolderList, setIsLoadingFolderList] = useState(false);
@@ -82,6 +83,8 @@ const RemoteSyncConfigurator: FC = () => {
             if (response.status === 200) {
                 queryClient.clear();
                 setReportLocation('remote');
+                setSelectedDevice(0);
+
                 navigate(ROUTES.OPERATIONS);
             }
         }

diff --git a/src/components/tensor-sharding-visualization/TensorVisualisationComponent.tsx b/src/components/tensor-sharding-visualization/TensorVisualisationComponent.tsx
@@ -56,6 +56,7 @@ const TensorVisualisationComponent: React.FC<TensorVisualisationComponentProps>
 
     const [selectedTensix, setSelectedTensix] = useState<number | null>(null);
     const [chartData, setChartData] = useState<Partial<PlotData>[]>([]);
+
     if (!data || !devices) {
         return (
             <span className='tensor-visualisation-loader'>

diff --git a/src/hooks/useAPI.tsx b/src/hooks/useAPI.tsx
@@ -20,6 +20,7 @@ import {
 } from '../model/APIData';
 import { BufferType } from '../model/BufferType';
 import parseMemoryConfig, { MemoryConfig } from '../functions/parseMemoryConfig';
+import isValidNumber from '../functions/isValidNumber';
 
 export const fetchTabSession = async (): Promise<TabSession | null> => {
     // eslint-disable-next-line promise/valid-params
@@ -170,20 +171,25 @@ export const useOperationsList = (deviceId?: number) => {
     });
 };
 
-export const useOperationDetails = (operationId: number | null) => {
+export const useOperationDetails = (operationId: number | null, deviceId?: number | null) => {
     const { data: operations } = useOperationsList();
-    const operation = operations?.filter((_operation) => {
-        return _operation.id === operationId;
-    })[0];
+    const operation = operations?.filter((_operation) => _operation.id === operationId)[0];
+
     const operationDetails = useQuery<OperationDetailsData>(
-        ['get-operation-detail', operationId],
+        ['get-operation-detail', operationId, deviceId],
         () => fetchOperationDetails(operationId),
         {
             retry: 2,
             retryDelay: (retryAttempt) => Math.min(retryAttempt * 100, 500),
         },
     );
 
+    if (operationDetails.data) {
+        operationDetails.data.buffers = operationDetails.data.buffers.filter((buffer) =>
+            isValidNumber(deviceId) ? buffer.device_id === deviceId : true,
+        );
+    }
+
     return {
         operation,
         operationDetails,

diff --git a/src/scss/components/BufferSummaryPlot.scss b/src/scss/components/BufferSummaryPlot.scss
@@ -11,6 +11,13 @@ $height-offset: 15px;
 .buffer-summary-chart {
     margin-bottom: $height-offset * 2;
 
+    .controls {
+        display: inline-flex;
+        flex-direction: column;
+        margin-bottom: 0;
+        align-items: flex-start;
+    }
+
     .chart-position {
         position: relative;
     }

diff --git a/src/scss/components/OperationDetailsComponent.scss b/src/scss/components/OperationDetailsComponent.scss
@@ -16,8 +16,11 @@
         }
     }
 
-    .zoom-controls {
+    .controls {
+        display: inline-flex;
+        flex-direction: column;
         margin-bottom: 0;
+        align-items: flex-start;
     }
 
     .circular-buffers-plot-title {
@@ -110,8 +113,8 @@
     }
 
     .no-buffer-type-selected {
-        margin-top: 20px;
-        margin-bottom: 0;
+        margin-top: 30px;
+        margin-bottom: 30px;
     }
 
     .deallocate-msg {

diff --git a/src/store/app.ts b/src/store/app.ts
@@ -31,4 +31,4 @@ export const fileTransferProgressAtom = atom({
     status: FileStatus.FINISHED,
 });
 
-export const selectedDeviceAtom = atom<number>();
+export const selectedDeviceAtom = atom<number | null>(0); // Assumes device_id always uses a zero based index