Commit f762d98
authored
[Fix, NPU] Set soc_info for the NPU device (#1149)
## Summary
With the `teb` update, the `soc version` must be explicitly set in
`set_current_compile_soc_info` for `get_soc_spec` to return the correct
`UB_SIZE`.
```
import torch
from tbe.common.platform import get_soc_spec
from tbe.common.platform import set_current_compile_soc_info
ub_size = get_soc_spec("UB_SIZE")
print(f"UB size: {ub_size}")
```
<img width="646" height="32" alt="image"
src="https://github.com/user-attachments/assets/d57e52e5-e468-4e46-a4db-94351364d218"
/> ❌
## Testing Done
### Atlas 900 A2 PoD
#### Script: ub_size_test.py
```python
import torch
from tbe.common.platform import get_soc_spec
from tbe.common.platform import set_current_compile_soc_info
device = getattr(torch, "npu")
soc_info = device.get_device_name(device.current_device())
print("Current SOC info:", soc_info)
ub_size = get_soc_spec("UB_SIZE")
print(f"UB size: {ub_size}")
```
#### Result
<img width="713" height="304" alt="image"
src="https://github.com/user-attachments/assets/56215e32-1c1b-49ed-9e5b-a859e80729d1"
/>
### Atlas 800I A2
Thanks @TianHao324 for verifying on the Atlas 800I A2 — results are
correct there too.
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence1 parent 05f5878 commit f762d98
1 file changed
+6
-0
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
198 | 198 | | |
199 | 199 | | |
200 | 200 | | |
| 201 | + | |
| 202 | + | |
| 203 | + | |
| 204 | + | |
| 205 | + | |
| 206 | + | |
201 | 207 | | |
202 | 208 | | |
203 | 209 | | |
| |||
0 commit comments