install_fbgemm.md

这份 `README.md` 是专门为你提供的 **海光 DCU (DTK 25.04.2)** 环境定制的。它整合了你提供的 Docker 镜像信息、编译步骤，以及我们之前解决的 **OneTBB 路径修复**和 **HIP 类型适配**等关键补丁。

---

# FBGEMM_GPU 安装指南 (DTK 25.04.2)

本指南用于在海光加速卡环境下，基于 **DTK 25.04.2** 和 **PyTorch 2.5.1/2.7** 编译并安装 `fbgemm_gpu`。

## 1. 环境启动

使用包含 DTK 25.04.2 的 PyTorch 镜像启动容器：

```bash
# 拉取镜像
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.2-py3.10

# 启动容器
docker run -it \
  --network=host \
  --ipc=host \
  --shm-size=16G \
  --device=/dev/kfd \
  --device=/dev/mkfd \
  --device=/dev/dri \
  -v /opt/hyhal:/opt/hyhal \
  --group-add video \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.2-py3.10

```

## 2. 获取源码

下载并初始化 `FBGEMM` 仓库（推荐版本 v1.3.0）：

```bash
git clone https://github.com/pytorch/FBGEMM.git --branch v1.3.0
cd FBGEMM/fbgemm_gpu
git submodule update --init --recursive


```

## 3. 关键补丁与修复 (必须执行)

在编译前，需要针对 DTK 25.04.2 的类型定义和 TBB 库结构进行手动修复：
### a. 系统环境与代码清理

```bash
# 建立系统库链接
ln -s /lib/x86_64-linux-gnu/librt.so.1 /usr/lib/x86_64-linux-gnu/librt.so

```
### b. 移除断言检查  CUDA_KERNEL_ASSERT2 
![cuda_kernel_assert2](./imgs/cuda_kernel_assert2.png)

### c. 移除 ./include/fbgemm_gpu/utils/kernel_launcher.cuh 中的 TORCH_CHECK
![torch_check](./imgs/torch_check.png)


## 4. 编译与安装
```bash
# 1. 清理缓存
rm -rf _skbuild build

# 2. 设置环境变量
export PYTORCH_ROCM_ARCH="gfx936"
export ROCM_PATH=/opt/dtk
export BUILD_VERSION="1.3.0"
export TORCH_PATH="/usr/local/lib/python3.10/dist-packages/torch/share/cmake/Torch"

# 3. 执行安装命令
# 注意：通过 --no-as-needed 强制链接 TBB 并烧录 RPATH 以免运行时报错
python setup.py bdist_wheel \
  --package_channel=release \
  -DFBGEMM_BUILD_TARGET=default \
  -DFBGEMM_BUILD_VARIANT=rocm \
  -DCMAKE_PREFIX_PATH="${TORCH_PATH};/opt/dtk" \
  -DAMDGPU_TARGETS="${PYTORCH_ROCM_ARCH}" \
  -DTORCH_DIR="${TORCH_PATH}" \
  -DCMAKE_C_FLAGS="-Wno-return-type -Wno-ignored-attributes -O1" \
  -DCMAKE_CXX_FLAGS="-Wno-return-type -Wno-ignored-attributes -O1 -DhipDeviceProp_t=hipDeviceProp_t_v2" \
  -DCMAKE_HIP_FLAGS="-Wno-return-type -O1 -DhipDeviceProp_t=hipDeviceProp_t_v2" \
  -DUSE_AVX512=on \
  -DCOPY_VISIBLE_LIBRARIES=ON \
  -DCMAKE_SHARED_LINKER_FLAGS="-Wl,--no-as-needed -ltbb -Wl,-rpath,/usr/lib/x86_64-linux-gnu" \
  -DTBB_INCLUDE_DIR=/usr/include

```

## 5. 验证安装

安装完成后，在 Python 中验证：

```python
import torch
import fbgemm_gpu
print("FBGEMM_GPU 导入成功！")

```