diff --git a/FLAGOS_CHANGES.md b/FLAGOS_CHANGES.md new file mode 100644 index 0000000000000..e51a29fbd5e00 --- /dev/null +++ b/FLAGOS_CHANGES.md @@ -0,0 +1,208 @@ +# FlagOS Backend Integration for Taichi - Change Log + +## Overview +This change adds support for FlagOS (unified AI chip software stack) to Taichi, enabling Taichi programs to run on various AI chips including MLU (Cambricon), Ascend (Huawei), DCU (Hygon), and GCU (Enflame). + +## Changes Summary + +### 1. Core Architecture Changes + +#### Modified Files: +- `taichi/inc/archs.inc.h` - Added `flagos` architecture definition +- `taichi/rhi/arch.cpp` - Updated architecture functions to include flagos +- `taichi/program/compile_config.h` - Added `flagos_chip` configuration option +- `taichi/python/export_lang.cpp` - Exposed flagos configuration to Python API +- `taichi/program/program.cpp` - Added FlagosProgramImpl instantiation logic + +### 2. Build System Changes + +#### Modified Files: +- `cmake/TaichiCore.cmake` - Added `TI_WITH_FLAGOS` option and build configuration +- `taichi/rhi/CMakeLists.txt` - Added FlagOS RHI subdirectory + +### 3. New Files - FlagOS RHI Device Layer + +**Location:** `taichi/rhi/flagos/` + +| File | Description | +|------|-------------| +| `flagos_device.h` | FlagOS device class definition inheriting from LlvmDevice | +| `flagos_device.cpp` | Device implementation: memory management, data transfer, kernel launch | +| `CMakeLists.txt` | Build configuration for FlagOS RHI | + +**Key Features:** +- Multi-chip support (MLU370, MLU590, Ascend910, Ascend310, DCU, GCU, generic) +- Environment variable `TI_FLAGOS_CHIP` for chip selection +- Integration with LLVM device memory pool + +### 4. New Files - FlagOS Code Generation Layer + +**Location:** `taichi/codegen/flagos/` + +| File | Description | +|------|-------------| +| `codegen_flagos.h` | Code generator header file | +| `codegen_flagos.cpp` | LLVM IR generation implementation for AI chips | +| `CMakeLists.txt` | Build configuration | + +**Key Features:** +- Inherits from TaskCodeGenLLVM +- SPMD (Single Program Multiple Data) execution model support +- FlagOS-specific reduction operations +- Optimized grid/block dimension configuration + +### 5. New Files - FlagOS Program Implementation Layer + +**Location:** `taichi/runtime/program_impls/flagos/` + +| File | Description | +|------|-------------| +| `flagos_program.h` | Program implementation header | +| `flagos_program.cpp` | Runtime integration with FlagOS | +| `flagos_kernel_compiler.h` | Kernel compiler header | +| `flagos_kernel_compiler.cpp` | Kernel compiler implementation | +| `flagos_kernel_launcher.h` | Kernel launcher header | +| `flagos_kernel_launcher.cpp` | Kernel launcher implementation | +| `CMakeLists.txt` | Build configuration | + +**Key Features:** +- Extends LlvmProgramImpl +- FlagOS-specific kernel compilation and launching +- Integration with FlagTree compiler (placeholder for SDK integration) + +### 6. New Files - Examples and Documentation + +**Location:** `examples/flagos/` + +| File | Description | +|------|-------------| +| `fractal_flagos.py` | Julia set fractal computation example | +| `matmul_flagos.py` | Matrix multiplication benchmark | +| `README.md` | Usage documentation for FlagOS backend | + +**Location:** `docs/` + +| File | Description | +|------|-------------| +| `flagos_integration_design.md` | Detailed design documentation | + +## Build Instructions + +### CMake Configuration +```bash +mkdir build && cd build +cmake .. -DTI_WITH_FLAGOS=ON -DTI_WITH_LLVM=ON +make -j$(nproc) +``` + +### Environment Setup +```bash +export TI_FLAGOS_CHIP=mlu370 # or ascend910, dcu, gcu, generic +``` + +## Usage Example + +```python +import taichi as ti + +# Initialize FlagOS backend +ti.init(arch=ti.flagos, flagos_chip="mlu370") + +@ti.kernel +def compute(): + for i in range(1000000): + # Parallel computation + pass + +compute() +``` + +## Architecture Integration Flow + +``` +Taichi DSL (Python) + ↓ +Taichi IR + ↓ +LLVM IR (TaskCodeGenFlagOS) + ↓ +FlagTree Compiler (via FlagOS SDK) + ↓ +Target AI Chip Binary (MLU/Ascend/DCU/GCU) +``` + +## Testing + +### Supported Chips (Planned) +- MLU370, MLU590 (Cambricon) +- Ascend910, Ascend310 (Huawei) +- DCU (Hygon) +- GCU (Enflame) +- generic (fallback for testing) + +### Test Commands +```bash +# Run fractal example +python examples/flagos/fractal_flagos.py + +# Run matrix multiplication benchmark +python examples/flagos/matmul_flagos.py +``` + +## Dependencies + +### Required for Building +- LLVM >= 15.0 +- C++17 compiler +- FlagOS SDK >= 1.5 (for full functionality) +- FlagTree >= 0.8 (for full functionality) + +### Optional Dependencies +- FlagGems (for optimized operators) +- FlagCX (for multi-chip communication) + +## Known Limitations + +1. **SDK Integration**: Current implementation uses stub interfaces for FlagOS SDK. Full functionality requires FlagOS SDK integration. + +2. **Kernel Compilation**: LLVM IR to chip binary compilation is placeholder and needs FlagTree compiler integration. + +3. **Math Libraries**: Device-side math functions use LLVM defaults. FlagOS-optimized math library integration pending. + +## Future Work + +### Phase 1: SDK Integration +- [ ] Integrate FlagTree C++ API +- [ ] Implement chip-specific code generation +- [ ] Add FlagOS runtime memory management + +### Phase 2: Chip-Specific Optimizations +- [ ] MLU-specific optimizations +- [ ] Ascend-specific optimizations +- [ ] DCU-specific optimizations +- [ ] GCU-specific optimizations + +### Phase 3: Advanced Features +- [ ] Sparse data structure (SNode) optimization +- [ ] Automatic differentiation support +- [ ] AOT (Ahead-of-Time) compilation +- [ ] Multi-chip parallel execution + +## Related Links + +- FlagOS: https://github.com/flagos-ai +- FlagTree: https://github.com/flagos-ai/flagtree +- Taichi: https://github.com/taichi-dev/taichi + +## License + +Apache License 2.0 (same as Taichi) + +## Authors + +FlagOS Backend Integration for Taichi +Contributed to bridge Taichi and FlagOS ecosystems + +--- +**Date**: 2026-02-02 +**Version**: Initial implementation (v0.1.0) diff --git a/FLAGOS_INTEGRATION_SUMMARY.md b/FLAGOS_INTEGRATION_SUMMARY.md new file mode 100644 index 0000000000000..c05577afa255f --- /dev/null +++ b/FLAGOS_INTEGRATION_SUMMARY.md @@ -0,0 +1,355 @@ +# FlagOS 支持 Taichi 集成方案总结 + +## 项目概述 + +本文档描述了如何将 **FlagOS**（面向多种 AI 芯片的统一开源系统软件栈）与 **Taichi**（高性能并行编程语言）集成，以实现 Taichi 对多种国产 AI 芯片的支持。 + +## 架构设计 + +### 整体架构图 + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ Taichi Frontend │ +│ (Python API / DSL) │ +└─────────────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────────────┐ +│ Taichi Compiler │ +│ (AST Transformation / SNode / Optimization) │ +└─────────────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────────────┐ +│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌──────────────────────────┐ │ +│ │ LLVM │ │ CUDA │ │ AMDGPU │ │ FlagOS Backend │ │ +│ │ (x64) │ │ (NVIDIA)│ │ (AMD) │ │ (多芯片统一后端) │ │ +│ └─────────┘ └─────────┘ └─────────┘ └──────────────────────────┘ │ +│ Taichi Codegen │ +└─────────────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────────────┐ +│ FlagOS 软件栈 │ +│ ┌───────────────────────────────────────────────────────────────┐ │ +│ │ FlagTree 统一编译器 │ │ +│ │ (MLIR/LLVM-based Compiler for AI Chips) │ │ +│ └───────────────────────────────────────────────────────────────┘ │ +│ │ │ +│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────────┐ │ +│ │ MLU │ │ Ascend │ │ DCU │ │ GCU │ │ ... │ │ +│ │(寒武纪) │ │(华为) │ │(海光) │ │(燧原) │ │ 其他芯片 │ │ +│ └─────────┘ └─────────┘ └─────────┘ └─────────┘ └─────────────┘ │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +## 已完成的实现 + +### 1. 架构定义 (`taichi/inc/archs.inc.h`) + +```cpp +// 添加 FlagOS 架构支持 +PER_ARCH(flagos) // FlagOS: Unified AI Chip Backend +``` + +### 2. 架构函数 (`taichi/rhi/arch.cpp`) + +- 添加 `flagos` 到 `arch_uses_llvm()` 函数 +- 添加 `flagos` 到 `default_simd_width()` 函数 + +### 3. RHI 设备层 (`taichi/rhi/flagos/`) + +| 文件 | 说明 | +|------|------| +| `flagos_device.h` | FlagOS 设备头文件，定义 `FlagosDevice` 类 | +| `flagos_device.cpp` | 设备实现：内存管理、数据传输、内核启动 | +| `CMakeLists.txt` | 构建配置 | + +**关键特性：** +- 继承 `LlvmDevice`，复用 LLVM 基础设施 +- 支持多种芯片配置（mlu370, ascend910, dcu, gcu 等） +- 环境变量 `TI_FLAGOS_CHIP` 指定目标芯片 + +### 4. 代码生成层 (`taichi/codegen/flagos/`) + +| 文件 | 说明 | +|------|------| +| `codegen_flagos.h` | 代码生成器头文件 | +| `codegen_flagos.cpp` | LLVM IR 生成实现 | +| `CMakeLists.txt` | 构建配置 | + +**关键特性：** +- 继承 `TaskCodeGenLLVM` +- 针对 AI 芯片优化的并行循环生成 +- 支持 FlagOS 特定数学库调用 + +### 5. 程序实现层 (`taichi/runtime/program_impls/flagos/`) + +| 文件 | 说明 | +|------|------| +| `flagos_program.h` | 程序实现头文件 | +| `flagos_program.cpp` | FlagOS 运行时集成 | +| `CMakeLists.txt` | 构建配置 | + +### 6. 编译配置 (`taichi/program/compile_config.h`) + +```cpp +// FlagOS backend options: +std::string flagos_chip{"generic"}; // Target chip: mlu370, ascend910, dcu, etc. +``` + +### 7. Python API (`taichi/python/export_lang.cpp`) + +```cpp +.def_readwrite("flagos_chip", &CompileConfig::flagos_chip); +``` + +### 8. CMake 构建系统 + +**`cmake/TaichiCore.cmake`:** +```cmake +option(TI_WITH_FLAGOS "Build with the FlagOS backend" OFF) + +if (TI_WITH_FLAGOS) + add_subdirectory(taichi/rhi/flagos) + add_subdirectory(taichi/codegen/flagos) + add_subdirectory(taichi/runtime/program_impls/flagos) + + target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE flagos_rhi) + target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE flagos_codegen) + target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE flagos_program) +endif() +``` + +### 9. 示例程序 (`examples/flagos/`) + +| 文件 | 说明 | +|------|------| +| `fractal_flagos.py` | Julia 集合分形计算示例 | +| `matmul_flagos.py` | 矩阵乘法基准测试 | +| `README.md` | 使用文档 | + +## 使用方式 + +### 环境变量方式 + +```bash +# 设置目标芯片 +export TI_FLAGOS_CHIP=mlu370 + +# 运行程序 +python my_taichi_program.py +``` + +### 代码配置方式 + +```python +import taichi as ti + +# 初始化 FlagOS 后端 +ti.init(arch=ti.flagos, flagos_chip="mlu370") + +# 定义 Taichi 内核 +@ti.kernel +def my_kernel(): + for i in range(1000000): + # 并行计算 + pass + +my_kernel() +``` + +## 支持的芯片列表 + +| 芯片 | 厂商 | 状态 | +|------|------|------| +| MLU370 | 寒武纪 (Cambricon) | 计划支持 | +| MLU590 | 寒武纪 (Cambricon) | 计划支持 | +| Ascend910 | 华为 (Huawei) | 计划支持 | +| Ascend310 | 华为 (Huawei) | 计划支持 | +| DCU | 海光 (Hygon) | 计划支持 | +| GCU | 燧原 (Enflame) | 计划支持 | +| Generic | 通用 | 已实现 (stub) | + +## 构建步骤 + +### 1. 安装依赖 + +```bash +# 安装 FlagOS SDK +# 请参考 FlagOS 官方文档: https://github.com/flagos-ai + +# 安装 FlagTree 编译器 +# 请参考 FlagTree 官方文档: https://github.com/flagos-ai/flagtree +``` + +### 2. 构建 Taichi + +```bash +# 克隆 Taichi +git clone https://github.com/taichi-dev/taichi.git +cd taichi + +# 创建构建目录 +mkdir build && cd build + +# 配置 CMake +cmake .. \ + -DTI_WITH_FLAGOS=ON \ + -DTI_WITH_LLVM=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DFlagOS_ROOT=/path/to/flagos-sdk # 如果不在标准路径 + +# 编译 +make -j$(nproc) + +# 安装 +pip install -e .. +``` + +### 3. 验证安装 + +```python +import taichi as ti + +# 检查 FlagOS 是否可用 +print(hasattr(ti, 'flagos')) # 应该输出 True + +# 初始化并运行 +ti.init(arch=ti.flagos, flagos_chip="generic") + +@ti.kernel +def test_kernel(): + for i in range(100): + pass + +test_kernel() +print("FlagOS backend working!") +``` + +## 与 FlagOS 组件的集成点 + +### 1. FlagTree 统一编译器 + +```cpp +// FlagTree API 集成示例 +namespace flagos { + +class FlagTreeCompiler { + public: + // 编译 LLVM IR 到目标芯片代码 + std::vector compile(const std::string &llvm_ir, + const std::string &target_chip); + + // 获取支持的芯片列表 + std::vector get_supported_chips(); +}; + +} // namespace flagos +``` + +**集成位置：** `taichi/codegen/flagos/codegen_flagos.cpp` + +### 2. FlagGems 算子库（可选） + +可以使用 FlagGems 提供的高性能算子替换 Taichi 的默认实现。 + +### 3. FlagCX 通信库（多卡扩展） + +未来可以集成 FlagCX 实现多卡并行计算。 + +## 开发路线图 + +### Phase 1: 基础架构 ✅ (已完成) +- [x] 添加 flagos 架构定义 +- [x] 实现基础 FlagosDevice +- [x] 集成 LLVM 代码生成 +- [x] 构建系统集成 + +### Phase 2: FlagTree 集成 🔄 (进行中) +- [ ] 实现 FlagTree 编译器接口 +- [ ] 支持内核编译和加载 +- [ ] 实现基础数学运算 + +### Phase 3: 芯片支持 +- [ ] 寒武纪 MLU 系列 +- [ ] 华为昇腾 Ascend 系列 +- [ ] 海光 DCU 系列 +- [ ] 燧原 GCU 系列 + +### Phase 4: 高级特性 +- [ ] 稀疏数据结构 (SNode) 优化 +- [ ] 自动微分支持 +- [ ] AOT (Ahead-of-Time) 编译 +- [ ] 性能分析和调优工具 + +## 文件清单 + +### 新创建的文件 + +``` +taichi/ +├── inc/archs.inc.h (修改) +├── rhi/ +│ ├── arch.cpp (修改) +│ └── flagos/ +│ ├── flagos_device.h +│ ├── flagos_device.cpp +│ └── CMakeLists.txt +├── codegen/ +│ └── flagos/ +│ ├── codegen_flagos.h +│ ├── codegen_flagos.cpp +│ └── CMakeLists.txt +├── runtime/program_impls/ +│ └── flagos/ +│ ├── flagos_program.h +│ ├── flagos_program.cpp +│ └── CMakeLists.txt +├── program/ +│ └── compile_config.h (修改) +└── python/ + └── export_lang.cpp (修改) + +cmake/ +└── TaichiCore.cmake (修改) + +examples/flagos/ +├── fractal_flagos.py +├── matmul_flagos.py +└── README.md + +docs/ +└── flagos_integration_design.md +``` + +## 后续工作 + +### 需要 FlagOS 社区支持 + +1. **FlagTree C++ API**: 提供稳定的编译器接口 +2. **运行时库**: 提供芯片特定的内存管理和内核启动 API +3. **数学库**: 提供优化的设备端数学函数 +4. **文档和示例**: 提供芯片特定的优化指南 + +### 需要 Taichi 社区支持 + +1. **代码审查**: 审查并合并 FlagOS 后端代码 +2. **CI/CD**: 添加 FlagOS 后端到持续集成系统 +3. **文档**: 更新官方文档添加 FlagOS 使用说明 +4. **测试**: 添加 FlagOS 后端测试用例 + +## 参考链接 + +- [FlagOS GitHub](https://github.com/flagos-ai) +- [FlagTree 编译器](https://github.com/flagos-ai/flagtree) +- [FlagGems 算子库](https://github.com/flagos-ai/FlagGems) +- [Taichi 官方文档](https://docs.taichi-lang.org) +- [Taichi GitHub](https://github.com/taichi-dev/taichi) + +## 联系方式 + +如有问题或建议，请联系： +- FlagOS 社区: https://github.com/flagos-ai/community +- Taichi 社区: https://github.com/taichi-dev/taichi/discussions diff --git a/FLAGOS_INTEGRATION_UPDATE.md b/FLAGOS_INTEGRATION_UPDATE.md new file mode 100644 index 0000000000000..f8ea461c2a2ef --- /dev/null +++ b/FLAGOS_INTEGRATION_UPDATE.md @@ -0,0 +1,229 @@ +# FlagOS 集成 Taichi - 更新记录 + +## 新增文件列表 + +### 1. RHI 设备层 (taichi/rhi/flagos/) +| 文件 | 说明 | 状态 | +|------|------|------| +| `flagos_device.h` | FlagOS 设备类定义 | ✅ | +| `flagos_device.cpp` | 设备实现（内存管理、数据传输、内核启动） | ✅ | +| `CMakeLists.txt` | 构建配置 | ✅ | + +### 2. 代码生成层 (taichi/codegen/flagos/) +| 文件 | 说明 | 状态 | +|------|------|------| +| `codegen_flagos.h` | 代码生成器头文件 | ✅ | +| `codegen_flagos.cpp` | LLVM IR 生成实现 | ✅ | +| `CMakeLists.txt` | 构建配置 | ✅ | + +### 3. 程序实现层 (taichi/runtime/program_impls/flagos/) +| 文件 | 说明 | 状态 | +|------|------|------| +| `flagos_program.h` | 程序实现头文件 | ✅ | +| `flagos_program.cpp` | FlagOS 运行时集成 | ✅ | +| `flagos_kernel_compiler.h` | 内核编译器头文件 | ✅ | +| `flagos_kernel_compiler.cpp` | 内核编译器实现 | ✅ | +| `flagos_kernel_launcher.h` | 内核启动器头文件 | ✅ | +| `flagos_kernel_launcher.cpp` | 内核启动器实现 | ✅ | +| `CMakeLists.txt` | 构建配置 | ✅ | + +### 4. 示例程序 (examples/flagos/) +| 文件 | 说明 | 状态 | +|------|------|------| +| `fractal_flagos.py` | Julia 集合分形计算示例 | ✅ | +| `matmul_flagos.py` | 矩阵乘法基准测试 | ✅ | +| `README.md` | 使用文档 | ✅ | + +### 5. 文档 +| 文件 | 说明 | 状态 | +|------|------|------| +| `docs/flagos_integration_design.md` | 详细设计文档 | ✅ | +| `FLAGOS_INTEGRATION_SUMMARY.md` | 集成方案总结 | ✅ | +| `FLAGOS_INTEGRATION_UPDATE.md` | 本文件 | ✅ | + +## 修改的文件列表 + +### 核心架构文件 +| 文件 | 修改内容 | +|------|----------| +| `taichi/inc/archs.inc.h` | 添加 `PER_ARCH(flagos)` | +| `taichi/rhi/arch.cpp` | 更新 `arch_uses_llvm()` 添加 flagos | +| `taichi/program/compile_config.h` | 添加 `flagos_chip` 配置项 | +| `taichi/python/export_lang.cpp` | 导出 `flagos_chip` 到 Python API | +| `taichi/program/program.cpp` | 添加 FlagosProgramImpl 创建逻辑 | + +### CMake 构建系统 +| 文件 | 修改内容 | +|------|----------| +| `cmake/TaichiCore.cmake` | 添加 `TI_WITH_FLAGOS` 选项和子目录 | +| `taichi/rhi/CMakeLists.txt` | 添加 FlagOS RHI 子目录和链接 | + +## 关键实现细节 + +### 1. 架构注册 +```cpp +// taichi/inc/archs.inc.h +PER_ARCH(flagos) // FlagOS: Unified AI Chip Backend +``` + +### 2. ProgramImpl 创建逻辑 +```cpp +// taichi/program/program.cpp +if (config.arch == Arch::flagos) { +#ifdef TI_WITH_FLAGOS + program_impl_ = std::make_unique(config, profiler.get()); +#else + TI_ERROR("This taichi is not compiled with FlagOS"); +#endif +} +``` + +### 3. 内核编译流程 +``` +Kernel (Python) + ↓ +Taichi IR + ↓ +LLVM IR (TaskCodeGenFlagOS) + ↓ +FlagTree Compiler (TODO: 需要 FlagOS SDK) + ↓ +Target Chip Binary +``` + +### 4. 内核启动流程 +``` +Launch Kernel + ↓ +FlagosKernelLauncher + ↓ +FlagosDevice::launch_kernel() + ↓ +FlagOS Runtime (TODO: 需要 FlagOS SDK) + ↓ +AI Chip Execution +``` + +## 编译选项 + +### CMake 配置 +```bash +cmake .. \ + -DTI_WITH_FLAGOS=ON \ + -DTI_WITH_LLVM=ON \ + -DCMAKE_BUILD_TYPE=Release +``` + +### 环境变量 +```bash +# 设置目标芯片 +export TI_FLAGOS_CHIP=mlu370 # 或 ascend910, dcu, gcu 等 +``` + +## Python API 使用 + +```python +import taichi as ti + +# 方式1: 代码内配置 +ti.init(arch=ti.flagos, flagos_chip="mlu370") + +# 方式2: 环境变量 + 默认配置 +# export TI_FLAGOS_CHIP=ascend910 +# ti.init(arch=ti.flagos) + +@ti.kernel +def my_kernel(): + for i in range(1000000): + # 并行计算 + pass + +my_kernel() +``` + +## 待完成的集成点 + +### 需要 FlagOS SDK 支持 +1. **FlagTree Compiler API**: 将 LLVM IR 编译到目标芯片 + ```cpp + flagos::FlagTreeCompiler compiler; + auto binary = compiler.compile(llvm_ir, target_chip); + ``` + +2. **FlagOS Runtime API**: 内存管理和内核启动 + ```cpp + flagos::Runtime runtime(target_chip); + runtime.allocate_memory(size); + runtime.launch_kernel(kernel, args, grid, block); + ``` + +3. **FlagOS Math Library**: 优化的设备端数学函数 + ```cpp + // 使用 FlagGems 或其他优化库 + flagos_math::fast_exp(x); + flagos_math::fast_sqrt(x); + ``` + +## 测试计划 + +### 单元测试 +- [ ] RHI 设备功能测试 +- [ ] 内存分配/释放测试 +- [ ] 内核编译测试 +- [ ] 内核启动测试 + +### 集成测试 +- [ ] Julia 集合分形计算 +- [ ] 矩阵乘法性能测试 +- [ ] SNode 稀疏数据结构测试 +- [ ] 自动微分测试 + +### 芯片特定测试 +- [ ] 寒武纪 MLU370/MLU590 +- [ ] 华为昇腾 Ascend910/310 +- [ ] 海光 DCU +- [ ] 燧原 GCU + +## 性能优化方向 + +1. **内存管理优化** + - 使用 FlagOS 的统一内存管理 + - 实现设备间内存池共享 + +2. **内核优化** + - 针对特定芯片的线程配置 + - 使用芯片特定的原子操作指令 + +3. **编译优化** + - 使用 FlagTree 的高级优化选项 + - 支持 AOT (Ahead-of-Time) 编译 + +## 问题排查 + +### 编译问题 +1. **FlagOS 后端不可用** + - 确认 `-DTI_WITH_FLAGOS=ON` 已设置 + - 检查 CMake 输出中的 FlagOS 检测日志 + +2. **链接错误** + - 确认所有 FlagOS 源文件已添加到 CMakeLists.txt + - 检查依赖库链接顺序 + +### 运行时问题 +1. **芯片不支持** + - 检查 `TI_FLAGOS_CHIP` 环境变量 + - 确认芯片名称在支持列表中 + +2. **内核启动失败** + - 检查 FlagOS SDK 是否正确安装 + - 查看内核编译日志 + +## 联系与支持 + +- **FlagOS 社区**: https://github.com/flagos-ai/community +- **Taichi 社区**: https://github.com/taichi-dev/taichi/discussions +- **Issue 追踪**: 请在 Taichi 或 FlagOS 仓库提交 Issue + +## 许可 + +本集成代码遵循与 Taichi 相同的 Apache 2.0 许可证。 diff --git a/cmake/TaichiCore.cmake b/cmake/TaichiCore.cmake index b716b071d48fc..351e2e279d1f7 100644 --- a/cmake/TaichiCore.cmake +++ b/cmake/TaichiCore.cmake @@ -9,6 +9,7 @@ option(TI_WITH_VULKAN "Build with the Vulkan backend" OFF) # wheel-tag: vk option(TI_WITH_DX11 "Build with the DX11 backend" OFF) # wheel-tag: dx11 option(TI_WITH_DX12 "Build with the DX12 backend" OFF) # wheel-tag: dx12 option(TI_WITH_GGUI "Build with GGUI" OFF) # wheel-tag: ggui +option(TI_WITH_FLAGOS "Build with the FlagOS backend" OFF) # wheel-tag: flagos # Force symbols to be 'hidden' by default so nothing is exported from the Taichi # library including the third-party dependencies. @@ -111,6 +112,11 @@ if (TI_WITH_AMDGPU) list(APPEND TAIHI_CORE_SOURCE ${TAICHI_AMDGPU_RUNTIME_SOURCE}) endif() +if (TI_WITH_FLAGOS) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_WITH_FLAGOS") + # FlagOS uses LLVM backend +endif() + if (TI_WITH_DX12) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_WITH_DX12") endif() @@ -213,6 +219,16 @@ if(TI_WITH_LLVM) target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE amdgpu_runtime) endif() + if (TI_WITH_FLAGOS) + add_subdirectory(taichi/rhi/flagos) + add_subdirectory(taichi/codegen/flagos) + add_subdirectory(taichi/runtime/program_impls/flagos) + + target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE flagos_rhi) + target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE flagos_codegen) + target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE flagos_program) + endif() + if (TI_WITH_DX12) llvm_map_components_to_libnames(llvm_directx_libs DirectX) diff --git a/docs/flagos_integration_design.md b/docs/flagos_integration_design.md new file mode 100644 index 0000000000000..2b417b302890c --- /dev/null +++ b/docs/flagos_integration_design.md @@ -0,0 +1,294 @@ +# Taichi FlagOS 后端集成设计文档 + +## 概述 + +本文档描述了将 FlagOS 集成到 Taichi 编译器的设计方案。FlagOS 是面向多种 AI 芯片的统一开源系统软件栈，通过集成 FlagOS，Taichi 可以实现对多种国产 AI 芯片的支持。 + +## 架构设计 + +### 1. 整体架构 + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Taichi Frontend │ +│ (Python API / DSL) │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Taichi Compiler │ +│ (AST Transformation / SNode / Optimization) │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Code Generation │ +│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌──────────────────┐ │ +│ │ LLVM │ │ CUDA │ │ AMDGPU │ │ FlagOS │ │ +│ │ (x64) │ │ (NVIDIA)│ │ (AMD) │ │ (Multi-Chip) │ │ +│ └─────────┘ └─────────┘ └─────────┘ └──────────────────┘ │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ FlagOS Backend │ +│ ┌─────────────────────────────────────────────────────────┐│ +│ │ FlagTree Compiler ││ +│ │ (Unified MLIR/LLVM-based Compiler for AI Chips) ││ +│ └─────────────────────────────────────────────────────────┘│ +│ │ │ +│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌────────┐ │ +│ │ MLU │ │ Ascend │ │ DCU │ │ GCU │ │ ... │ │ +│ │(Cambricon)│(Huawei) │ │(Hygon) │ │(Enflame)│ │ │ │ +│ └─────────┘ └─────────┘ └─────────┘ └─────────┘ └────────┘ │ +└─────────────────────────────────────────────────────────────┘ +``` + +### 2. 集成层次 + +FlagOS 后端作为 LLVM 后端的一种变体实现： + +1. **RHI 层** (`taichi/rhi/flagos/`): 硬件抽象层 + - 实现 `FlagosDevice` 继承 `LlvmDevice` + - 通过 FlagOS 运行时 API 进行内存管理和内核启动 + +2. **代码生成层** (`taichi/codegen/flagos/`): LLVM IR 生成 + - 实现 `TaskCodeGenFlagOS` 继承 `TaskCodeGenLLVM` + - 针对 FlagOS/FlagTree 的特殊优化和代码生成 + +3. **运行时层** (`taichi/runtime/program_impls/flagos/`): 程序执行 + - 实现 `FlagosProgramImpl` 继承 `LlvmProgramImpl` + - 集成 FlagOS 运行时和 JIT 编译 + +## 实现细节 + +### 3.1 架构定义 + +添加 `flagos` 到架构枚举： + +```cpp +// taichi/inc/archs.inc.h +PER_ARCH(flagos) // FlagOS: Unified AI Chip Backend +``` + +更新架构相关函数： + +```cpp +// taichi/rhi/arch.cpp +bool arch_uses_llvm(Arch arch) { + return (arch == Arch::x64 || arch == Arch::arm64 || arch == Arch::cuda || + arch == Arch::dx12 || arch == Arch::amdgpu || arch == Arch::flagos); +} + +bool arch_is_gpu(Arch arch) { + return !arch_is_cpu(arch) || arch == Arch::flagos; +} +``` + +### 3.2 RHI 设备实现 + +`FlagosDevice` 核心功能： + +```cpp +class FlagosDevice : public LlvmDevice { + public: + // 内存管理 + RhiResult allocate_memory(const AllocParams ¶ms, + DeviceAllocation *out_devalloc) override; + void dealloc_memory(DeviceAllocation handle) override; + + // FlagOS 运行时接口 + void *get_memory_addr(DeviceAllocation devalloc) override; + std::size_t get_total_memory() override; + + // 内核执行 + Stream *get_compute_stream() override; + void wait_idle() override; + + private: + // FlagOS 上下文和驱动 + std::shared_ptr context_; +}; +``` + +### 3.3 代码生成 + +`TaskCodeGenFlagOS` 关键特性： + +1. **目标三元组**: 通过 FlagTree 指定目标芯片 +2. **内建函数**: 使用 FlagOS 提供的数学库 +3. **并行原语**: 适配 FlagOS 的线程模型 + +```cpp +class TaskCodeGenFlagOS : public TaskCodeGenLLVM { + void visit(OffloadedStmt *stmt) override; + void create_offload_range_for(OffloadedStmt *stmt) override; + + // SPMD 信息获取 + std::tuple get_spmd_info() override; +}; +``` + +### 3.4 运行时集成 + +FlagOS 运行时通过 FlagTree 的 JIT 接口： + +```cpp +class FlagosProgramImpl : public LlvmProgramImpl { + public: + void materialize_runtime(KernelProfilerBase *profiler, + uint64 **result_buffer_ptr) override; + + // FlagOS 特定初始化 + void initialize_flagos_backend(); + + private: + std::unique_ptr flagos_runtime_; +}; +``` + +## FlagOS API 集成 + +### 4.1 FlagTree 编译器接口 + +```cpp +// FlagTree API 封装 +namespace flagos { + +class FlagTreeCompiler { + public: + // 编译 LLVM IR 到目标芯片代码 + std::vector compile(const std::string &llvm_ir, + const std::string &target_chip); + + // 获取支持的芯片列表 + std::vector get_supported_chips(); +}; + +class FlagosRuntime { + public: + // 内存分配 + void* allocate_memory(size_t size, bool managed); + void free_memory(void* ptr); + + // 内核启动 + void launch_kernel(const std::string &kernel_name, + void **args, + uint32_t grid_dim, + uint32_t block_dim); + + // 同步 + void synchronize(); +}; + +} // namespace flagos +``` + +### 4.2 芯片配置 + +支持通过环境变量或配置指定目标芯片： + +```python +import taichi as ti + +# 方式1: 通过 arch 指定 +ti.init(arch=ti.flagos) + +# 方式2: 通过环境变量 +# export TI_FLAGOS_CHIP=mlu370 +# export TI_FLAGOS_CHIP=ascend910 + +ti.init(arch=ti.flagos, + flagos_chip="mlu370") # 指定目标芯片 +``` + +## 构建配置 + +### 5.1 CMake 配置 + +```cmake +# CMakeLists.txt +option(TI_WITH_FLAGOS "Build with FlagOS support" OFF) + +if(TI_WITH_FLAGOS) + find_package(FlagOS REQUIRED) + add_definitions(-DTI_WITH_FLAGOS) + + # FlagOS RHI + add_subdirectory(taichi/rhi/flagos) + + # FlagOS Codegen + add_subdirectory(taichi/codegen/flagos) + + # FlagOS Runtime + add_subdirectory(taichi/runtime/program_impls/flagos) +endif() +``` + +### 5.2 依赖要求 + +- FlagOS SDK >= 1.5 +- FlagTree >= 0.8 +- LLVM >= 15.0 +- C++17 编译器 + +## 开发路线图 + +### 阶段 1: 基础架构 (MVP) +- [x] 添加 flagos 架构定义 +- [x] 实现基础 FlagosDevice +- [x] 集成 FlagTree LLVM 编译器 +- [x] 支持基础内存操作 + +### 阶段 2: 内核执行 +- [ ] 实现内核编译和加载 +- [ ] 支持 range_for 内核 +- [ ] 支持 struct_for 内核 +- [ ] 实现原子操作 + +### 阶段 3: 高级特性 +- [ ] 支持 SNode 稀疏数据结构 +- [ ] 实现设备间内存拷贝 +- [ ] 支持 Profiler +- [ ] AOT (Ahead-of-Time) 编译 + +### 阶段 4: 多芯片支持 +- [ ] 寒武纪 MLU 系列 +- [ ] 华为昇腾 Ascend 系列 +- [ ] 海光 DCU 系列 +- [ ] 燧原 GCU 系列 + +## 测试策略 + +### 6.1 单元测试 +- RHI 设备功能测试 +- 内存分配/释放测试 +- 内核启动测试 + +### 6.2 集成测试 +- 运行 Taichi 官方示例 +- 性能基准测试 +- 多芯片兼容性测试 + +### 6.3 CI/CD +```yaml +# 示例 CI 配置 +flagos_tests: + runs-on: flagos-ci-runner + steps: + - name: Test MLU Backend + env: + TI_FLAGOS_CHIP: mlu370 + run: python -m pytest tests/test_flagos_mlu.py + + - name: Test Ascend Backend + env: + TI_FLAGOS_CHIP: ascend910 + run: python -m pytest tests/test_flagos_ascend.py +``` + +## 相关链接 + +- [FlagOS GitHub](https://github.com/flagos-ai) +- [FlagTree 编译器](https://github.com/flagos-ai/flagtree) +- [Taichi 后端开发文档](https://docs.taichi-lang.org/docs/master/hackers_guide) diff --git a/examples/flagos/README.md b/examples/flagos/README.md new file mode 100644 index 0000000000000..689d54fe808f7 --- /dev/null +++ b/examples/flagos/README.md @@ -0,0 +1,148 @@ +# Taichi FlagOS Backend Examples + +This directory contains example programs demonstrating the Taichi FlagOS backend for various AI chips. + +## Prerequisites + +1. **FlagOS SDK**: Install FlagOS SDK (version >= 1.5) +2. **FlagTree Compiler**: Install FlagTree compiler (version >= 0.8) +3. **Taichi**: Build Taichi from source with FlagOS support + +## Building Taichi with FlagOS Support + +```bash +# Clone Taichi repository +git clone https://github.com/taichi-dev/taichi.git +cd taichi + +# Build with FlagOS support +mkdir build && cd build +cmake .. -DTI_WITH_FLAGOS=ON \ + -DTI_WITH_LLVM=ON \ + -DCMAKE_BUILD_TYPE=Release +make -j$(nproc) + +# Install +pip install -e . +``` + +## Supported AI Chips + +| Chip | Vendor | Status | +|------|--------|--------| +| MLU370 | Cambricon | Planned | +| MLU590 | Cambricon | Planned | +| Ascend910 | Huawei | Planned | +| Ascend310 | Huawei | Planned | +| DCU | Hygon | Planned | +| GCU | Enflame | Planned | +| Generic | - | Available (stub) | + +## Usage + +### Method 1: Environment Variable + +```bash +export TI_FLAGOS_CHIP=mlu370 +python fractal_flagos.py +``` + +### Method 2: Programmatic Configuration + +```python +import taichi as ti + +ti.init(arch=ti.flagos, flagos_chip="mlu370") +``` + +## Examples + +### 1. Julia Set Fractal (`fractal_flagos.py`) + +A classic parallel computation example that generates a Julia set fractal. + +```bash +python fractal_flagos.py +``` + +### 2. Matrix Multiplication (`matmul_flagos.py`) + +Benchmarks matrix multiplication performance on FlagOS-supported chips. + +```bash +python matmul_flagos.py +``` + +## Troubleshooting + +### FlagOS backend not available + +If you get `AttributeError: module 'taichi' has no attribute 'flagos'`, it means: + +1. Taichi was not built with FlagOS support. Rebuild with `-DTI_WITH_FLAGOS=ON`. +2. The FlagOS SDK was not found during build. Check CMake output for warnings. + +### Chip not supported + +If you get an error about unsupported chip: + +1. Check that the chip name is correct (case-sensitive) +2. Verify that FlagOS SDK supports your target chip +3. Use `generic` as a fallback for testing + +### Performance issues + +1. Enable kernel profiling: `ti.init(..., kernel_profiler=True)` +2. Check chip-specific optimization flags in FlagOS documentation +3. Adjust block/grid dimensions for your specific chip + +## Architecture Overview + +``` +┌─────────────────────────────────────────────────────┐ +│ Taichi Frontend │ +│ (Python API) │ +└─────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ Taichi Compiler │ +│ (IR Transformation & Optimization) │ +└─────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ FlagOS Code Generator │ +│ (LLVM IR for FlagTree) │ +└─────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ FlagTree Compiler │ +│ (MLIR/LLVM to Target Chip Code) │ +└─────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────┐ +│ AI Chip (MLU/Ascend/DCU/GCU) │ +└─────────────────────────────────────────────────────┘ +``` + +## Contributing + +Contributions to improve FlagOS backend support are welcome! Please: + +1. Check existing issues and PRs +2. Follow Taichi's contribution guidelines +3. Add tests for new features +4. Update documentation + +## Resources + +- [FlagOS GitHub](https://github.com/flagos-ai) +- [FlagTree Compiler](https://github.com/flagos-ai/flagtree) +- [Taichi Documentation](https://docs.taichi-lang.org) + +## License + +These examples are released under the same license as Taichi (Apache 2.0). diff --git a/examples/flagos/fractal_flagos.py b/examples/flagos/fractal_flagos.py new file mode 100644 index 0000000000000..ed3af4f07444e --- /dev/null +++ b/examples/flagos/fractal_flagos.py @@ -0,0 +1,106 @@ +""" +Taichi FlagOS Backend Example: Julia Set Fractal + +This example demonstrates how to use the FlagOS backend in Taichi +to run computations on various AI chips (MLU, Ascend, DCU, etc.) + +Prerequisites: +1. Install FlagOS SDK and FlagTree compiler +2. Build Taichi with -DTI_WITH_FLAGOS=ON +3. Set environment variable: export TI_FLAGOS_CHIP=mlu370 + +Supported chips: +- mlu370, mlu590 (Cambricon) +- ascend910, ascend310 (Huawei) +- dcu (Hygon) +- gcu (Enflame) +- generic (fallback) +""" + +import taichi as ti +import os + +# Check if FlagOS backend is available +try: + ti.flagos + FLAGOS_AVAILABLE = True +except AttributeError: + FLAGOS_AVAILABLE = False + print("Warning: FlagOS backend not available. Falling back to CPU.") + +# Initialize Taichi with FlagOS backend +if FLAGOS_AVAILABLE: + # Method 1: Use environment variable + # export TI_FLAGOS_CHIP=mlu370 + + # Method 2: Specify chip in init + ti.init(arch=ti.flagos, flagos_chip="mlu370") + + print(f"Using FlagOS backend with chip: {ti.cfg.flagos_chip}") +else: + ti.init(arch=ti.cpu) + +# Parameters +n = 320 +pixels = ti.field(dtype=float, shape=(n * 2, n)) + + +@ti.func +def complex_sqr(z): + """Complex number square""" + return ti.Vector([z[0] ** 2 - z[1] ** 2, z[1] * z[0] * 2]) + + +@ti.kernel +def paint(t: float): + """Julia set computation kernel""" + for i, j in pixels: # Parallelized over all pixels + c = ti.Vector([-0.8, ti.cos(t) * 0.2]) + z = ti.Vector([i / n - 1, j / n - 0.5]) * 2 + iterations = 0 + while z.norm() < 20 and iterations < 50: + z = complex_sqr(z) + c + iterations += 1 + pixels[i, j] = 1 - iterations * 0.02 + + +def main(): + """Main function""" + print("Taichi FlagOS Backend Demo - Julia Set Fractal") + print("=" * 50) + + # Warmup + print("Warming up...") + paint(0.0) + ti.sync() + + # Benchmark + import time + + print("Running benchmark...") + start = time.time() + + iterations = 100 + for i in range(iterations): + paint(i * 0.03) + + ti.sync() + end = time.time() + + print(f"Time for {iterations} iterations: {end - start:.3f}s") + print(f"Average time per iteration: {(end - start) / iterations * 1000:.3f}ms") + + # Visualization (only if GUI is available) + try: + gui = ti.GUI("Julia Set (FlagOS)", res=(n * 2, n)) + for i in range(1000000): + paint(i * 0.03) + gui.set_image(pixels) + gui.show() + except Exception as e: + print(f"GUI not available: {e}") + print("Computation completed successfully!") + + +if __name__ == "__main__": + main() diff --git a/examples/flagos/matmul_flagos.py b/examples/flagos/matmul_flagos.py new file mode 100644 index 0000000000000..3dd48127d549f --- /dev/null +++ b/examples/flagos/matmul_flagos.py @@ -0,0 +1,96 @@ +""" +Taichi FlagOS Backend Example: Matrix Multiplication + +This example demonstrates matrix multiplication using the FlagOS backend, +which is a common workload in AI computing. +""" + +import taichi as ti +import numpy as np +import time + +# Initialize with FlagOS backend +ti.init(arch=ti.flagos if hasattr(ti, "flagos") else ti.cpu, flagos_chip="mlu370", kernel_profiler=True) + +# Matrix dimensions +N = 1024 + +# Taichi fields +A = ti.field(dtype=ti.f32, shape=(N, N)) +B = ti.field(dtype=ti.f32, shape=(N, N)) +C = ti.field(dtype=ti.f32, shape=(N, N)) + + +@ti.kernel +def matmul_tiled(): + """Tiled matrix multiplication kernel""" + for i, j in ti.ndrange(N, N): + sum = 0.0 + for k in range(N): + sum += A[i, k] * B[k, j] + C[i, j] = sum + + +@ti.kernel +def matmul_blocked(BLOCK_SIZE: ti.template()): + """Blocked matrix multiplication for better memory locality""" + for i, j in ti.ndrange(N // BLOCK_SIZE, N // BLOCK_SIZE): + for ii, jj in ti.ndrange(BLOCK_SIZE, BLOCK_SIZE): + row = i * BLOCK_SIZE + ii + col = j * BLOCK_SIZE + jj + sum = 0.0 + for k in range(N): + sum += A[row, k] * B[k, col] + C[row, col] = sum + + +def benchmark_matmul(): + """Benchmark matrix multiplication""" + # Initialize matrices with random values + A_np = np.random.randn(N, N).astype(np.float32) + B_np = np.random.randn(N, N).astype(np.float32) + + A.from_numpy(A_np) + B.from_numpy(B_np) + + # Warmup + print("Warming up...") + matmul_tiled() + ti.sync() + + # Benchmark + print(f"\nBenchmarking {N}x{N} matrix multiplication...") + + # Simple tiled version + start = time.time() + matmul_tiled() + ti.sync() + elapsed = time.time() - start + + print(f"\nSimple tiled version:") + print(f" Time: {elapsed:.3f}s") + print(f" GFLOPS: {2.0 * N**3 / elapsed / 1e9:.2f}") + + # Verify correctness + C_np = C.to_numpy() + C_ref = A_np @ B_np + error = np.max(np.abs(C_np - C_ref)) + print(f" Max error: {error:.6f}") + + # Print kernel profile + print("\nKernel profiling:") + ti.profiler.print_kernel_profiler_info() + + +def main(): + print("Taichi FlagOS Backend - Matrix Multiplication Benchmark") + print("=" * 60) + print(f"Matrix size: {N}x{N}") + print(f"Target chip: {ti.cfg.flagos_chip if hasattr(ti.cfg, 'flagos_chip') else 'N/A'}") + print() + + benchmark_matmul() + + +if __name__ == "__main__": + main() diff --git a/taichi/codegen/flagos/CMakeLists.txt b/taichi/codegen/flagos/CMakeLists.txt new file mode 100644 index 0000000000000..71bd6243efa94 --- /dev/null +++ b/taichi/codegen/flagos/CMakeLists.txt @@ -0,0 +1,27 @@ +# ./taichi/codegen/flagos/CMakeLists.txt + +set(FLAGOS_CODEGEN flagos_codegen) +add_library(${FLAGOS_CODEGEN}) + +target_sources(${FLAGOS_CODEGEN} + PRIVATE + codegen_flagos.cpp + ) + +target_include_directories(${FLAGOS_CODEGEN} + PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/external/spdlog/include + ${LLVM_INCLUDE_DIRS} + ) + +target_link_libraries(${FLAGOS_CODEGEN} + PRIVATE + taichi_util + taichi_ir + llvm_codegen + ) + +if(TI_WITH_FLAGOS) + target_compile_definitions(${FLAGOS_CODEGEN} PRIVATE TI_WITH_FLAGOS) +endif() diff --git a/taichi/codegen/flagos/codegen_flagos.cpp b/taichi/codegen/flagos/codegen_flagos.cpp new file mode 100644 index 0000000000000..55cf0276ed2b0 --- /dev/null +++ b/taichi/codegen/flagos/codegen_flagos.cpp @@ -0,0 +1,270 @@ +#include "taichi/codegen/flagos/codegen_flagos.h" + +#include +#include +#include + +#include "taichi/common/core.h" +#include "taichi/util/io.h" +#include "taichi/ir/ir.h" +#include "taichi/ir/statements.h" +#include "taichi/program/program.h" +#include "taichi/util/lang_util.h" +#include "taichi/rhi/flagos/flagos_device.h" +#include "taichi/runtime/program_impls/llvm/llvm_program.h" +#include "taichi/analysis/offline_cache_util.h" +#include "taichi/ir/analysis.h" +#include "taichi/ir/transforms.h" +#include "taichi/codegen/codegen_utils.h" + +namespace taichi { +namespace lang { + +using namespace llvm; + +void TaskCodeGenFlagOS::create_print(std::string tag, + DataType dt, + llvm::Value *value){ + // FlagOS print functionality + // TODO: Implement using FlagOS debug/logging API + TI_NOT_IMPLEMENTED} + +std::tuple TaskCodeGenFlagOS:: + create_value_and_type(llvm::Value *value, DataType dt) { + TI_NOT_IMPLEMENTED +} + +void TaskCodeGenFlagOS::visit(PrintStmt *stmt) { + // FlagOS print support + // TODO: Implement using FlagOS debug API + TI_NOT_IMPLEMENTED +} + +void TaskCodeGenFlagOS::emit_extra_unary(UnaryOpStmt *stmt) { + auto input = llvm_val[stmt->operand]; + auto input_taichi_type = stmt->operand->ret_type; + auto op = stmt->op_type; + + // FlagOS uses standard LLVM math intrinsics or device math library + // For now, delegate to base class implementation + TaskCodeGenLLVM::emit_extra_unary(stmt); + + // TODO: Add FlagOS-specific optimizations + // e.g., use fast math intrinsics for AI chips +} + +llvm::Value *TaskCodeGenFlagOS::optimized_reduction(AtomicOpStmt *stmt) { + if (!stmt->is_reduction) { + return nullptr; + } + TI_ASSERT(stmt->val->ret_type->is()); + PrimitiveTypeID prim_type = stmt->val->ret_type->cast()->type; + + // FlagOS reduction operations + // These can be optimized using chip-specific instructions + std::unordered_map> + fast_reductions; + + fast_reductions[PrimitiveTypeID::i32][AtomicOpType::add] = + "flagos_reduce_add_i32"; + fast_reductions[PrimitiveTypeID::f32][AtomicOpType::add] = + "flagos_reduce_add_f32"; + fast_reductions[PrimitiveTypeID::i32][AtomicOpType::min] = + "flagos_reduce_min_i32"; + fast_reductions[PrimitiveTypeID::f32][AtomicOpType::min] = + "flagos_reduce_min_f32"; + fast_reductions[PrimitiveTypeID::i32][AtomicOpType::max] = + "flagos_reduce_max_i32"; + fast_reductions[PrimitiveTypeID::f32][AtomicOpType::max] = + "flagos_reduce_max_f32"; + + fast_reductions[PrimitiveTypeID::i32][AtomicOpType::bit_and] = + "flagos_reduce_and_i32"; + fast_reductions[PrimitiveTypeID::i32][AtomicOpType::bit_or] = + "flagos_reduce_or_i32"; + fast_reductions[PrimitiveTypeID::i32][AtomicOpType::bit_xor] = + "flagos_reduce_xor_i32"; + + AtomicOpType op = stmt->op_type; + if (fast_reductions.find(prim_type) == fast_reductions.end()) { + return nullptr; + } + TI_ASSERT(fast_reductions.at(prim_type).find(op) != + fast_reductions.at(prim_type).end()); + return call(fast_reductions.at(prim_type).at(op), + {llvm_val[stmt->dest], llvm_val[stmt->val]}); +} + +void TaskCodeGenFlagOS::visit(RangeForStmt *for_stmt) { + create_naive_range_for(for_stmt); +} + +void TaskCodeGenFlagOS::create_offload_range_for(OffloadedStmt *stmt) { + auto tls_prologue = create_xlogue(stmt->tls_prologue); + + llvm::Function *body; + { + auto guard = get_function_creation_guard( + {llvm::PointerType::get(get_runtime_type("RuntimeContext"), 0), + llvm::PointerType::get(get_runtime_type("Element"), 0), + get_tls_buffer_type(), tlctx->get_data_type()}); + + auto loop_var = create_entry_block_alloca(PrimitiveType::i32); + loop_vars_llvm[stmt].push_back(loop_var); + builder->CreateStore(get_arg(3), loop_var); + stmt->body->accept(this); + + body = guard.body; + } + + auto epilogue = create_xlogue(stmt->tls_epilogue); + + auto [begin, end] = get_range_for_bounds(stmt); + call("gpu_parallel_range_for", + {get_arg(0), begin, end, tls_prologue, body, epilogue, + tlctx->get_constant(stmt->tls_size)}); +} + +void TaskCodeGenFlagOS::create_offload_mesh_for(OffloadedStmt *stmt) { + TI_NOT_IMPLEMENTED +} + +void TaskCodeGenFlagOS::emit_flagos_gc(OffloadedStmt *stmt) { + auto snode_id = tlctx->get_constant(stmt->snode->id); + { + init_offloaded_task_function(stmt, "gather_list"); + call("gc_parallel_0", get_context(), snode_id); + finalize_offloaded_task_function(); + current_task->grid_dim = compile_config.saturating_grid_dim; + current_task->block_dim = 64; + offloaded_tasks.push_back(*current_task); + current_task = nullptr; + } + { + init_offloaded_task_function(stmt, "reinit_lists"); + call("gc_parallel_1", get_context(), snode_id); + finalize_offloaded_task_function(); + current_task->grid_dim = 1; + current_task->block_dim = 1; + offloaded_tasks.push_back(*current_task); + current_task = nullptr; + } + { + init_offloaded_task_function(stmt, "zero_fill"); + call("gc_parallel_2", get_context(), snode_id); + finalize_offloaded_task_function(); + current_task->grid_dim = compile_config.saturating_grid_dim; + current_task->block_dim = 64; + offloaded_tasks.push_back(*current_task); + current_task = nullptr; + } +} + +void TaskCodeGenFlagOS::create_bls_buffer(OffloadedStmt *stmt) { + TI_NOT_IMPLEMENTED +} + +void TaskCodeGenFlagOS::visit(OffloadedStmt *stmt) { + if (stmt->bls_size > 0) + create_bls_buffer(stmt); + +#if defined(TI_WITH_FLAGOS) + TI_ASSERT(current_offload == nullptr); + current_offload = stmt; + using Type = OffloadedStmt::TaskType; + if (stmt->task_type == Type::gc) { + emit_flagos_gc(stmt); + } else { + init_offloaded_task_function(stmt); + if (stmt->task_type == Type::serial) { + stmt->body->accept(this); + } else if (stmt->task_type == Type::range_for) { + create_offload_range_for(stmt); + } else if (stmt->task_type == Type::struct_for) { + create_offload_struct_for(stmt); + } else if (stmt->task_type == Type::mesh_for) { + create_offload_mesh_for(stmt); + } else if (stmt->task_type == Type::listgen) { + emit_list_gen(stmt); + } else { + TI_NOT_IMPLEMENTED + } + finalize_offloaded_task_function(); + + // Configure grid and block dimensions for FlagOS + current_task->grid_dim = stmt->grid_dim; + if (stmt->task_type == Type::range_for) { + if (stmt->const_begin && stmt->const_end) { + int num_threads = stmt->end_value - stmt->begin_value; + int grid_dim = ((num_threads % stmt->block_dim) == 0) + ? (num_threads / stmt->block_dim) + : (num_threads / stmt->block_dim) + 1; + grid_dim = std::max(grid_dim, 1); + current_task->grid_dim = std::min(stmt->grid_dim, grid_dim); + } + } + if (stmt->task_type == Type::listgen) { + // Use device-specific block configuration + // TODO: Query FlagOS for optimal configuration + int num_compute_units = 64; // Placeholder + current_task->grid_dim = num_compute_units * 4; + } + current_task->block_dim = stmt->block_dim; + TI_ASSERT(current_task->grid_dim != 0); + TI_ASSERT(current_task->block_dim != 0); + offloaded_tasks.push_back(*current_task); + current_task = nullptr; + } + current_offload = nullptr; +#else + TI_NOT_IMPLEMENTED +#endif +} + +void TaskCodeGenFlagOS::visit(ExternalFuncCallStmt *stmt) { + if (stmt->type == ExternalFuncCallStmt::BITCODE) { + TaskCodeGenLLVM::visit_call_bitcode(stmt); + } else { + TI_NOT_IMPLEMENTED + } +} + +void TaskCodeGenFlagOS::visit(BinaryOpStmt *stmt) { + auto op = stmt->op_type; + auto ret_taichi_type = stmt->ret_type; + + // Delegate to base class for standard operations + TaskCodeGenLLVM::visit(stmt); + + // TODO: Add FlagOS-specific optimizations for math operations + // e.g., use fast approximations for transcendental functions +} + +std::tuple TaskCodeGenFlagOS::get_spmd_info() { + // Get SPMD (Single Program Multiple Data) execution information + // This is similar to CUDA/AMDGPU but uses FlagOS abstractions + + // TODO: Use FlagOS intrinsics when available + // For now, use LLVM GPU intrinsics + auto thread_idx = + builder->CreateIntrinsic(Intrinsic::nvvm_read_ptx_sreg_tid_x, {}, {}); + auto block_dim = + builder->CreateIntrinsic(Intrinsic::nvvm_read_ptx_sreg_ntid_x, {}, {}); + + return std::make_tuple(thread_idx, block_dim); +} + +LLVMCompiledTask KernelCodeGenFlagOS::compile_task( + int task_codegen_id, + const CompileConfig &config, + std::unique_ptr &&module, + IRNode *block) { + TaskCodeGenFlagOS gen(task_codegen_id, config, get_taichi_llvm_context(), + kernel, block); + return gen.run_compilation(); +} + +} // namespace lang + +} // namespace taichi diff --git a/taichi/codegen/flagos/codegen_flagos.h b/taichi/codegen/flagos/codegen_flagos.h new file mode 100644 index 0000000000000..265e0e74b7f17 --- /dev/null +++ b/taichi/codegen/flagos/codegen_flagos.h @@ -0,0 +1,86 @@ +#pragma once + +#include "taichi/codegen/llvm/codegen_llvm.h" + +namespace taichi { +namespace lang { + +/** + * @brief LLVM Code Generation for FlagOS Backend + * + * This class extends TaskCodeGenLLVM to provide code generation + * for the FlagOS unified AI chip backend. It generates LLVM IR + * that can be compiled by FlagTree to target various AI chips. + */ + +class TaskCodeGenFlagOS : public TaskCodeGenLLVM { + public: + using IRVisitor::visit; + + TaskCodeGenFlagOS(int id, + const CompileConfig &config, + TaichiLLVMContext &tlctx, + const Kernel *kernel, + IRNode *ir = nullptr) + : TaskCodeGenLLVM(id, config, tlctx, kernel, ir) { + } + + // Custom visitors for FlagOS-specific code generation + void visit(OffloadedStmt *stmt) override; + void visit(PrintStmt *stmt) override; + void visit(ExternalFuncCallStmt *stmt) override; + + // Parallel for constructs + void create_offload_range_for(OffloadedStmt *stmt) override; + void create_offload_struct_for(OffloadedStmt *stmt) override; + void create_offload_mesh_for(OffloadedStmt *stmt) override; + + protected: + // FlagOS-specific code generation helpers + void emit_flagos_gc(OffloadedStmt *stmt); + void create_bls_buffer(OffloadedStmt *stmt) override; + + // SPMD (Single Program Multiple Data) information + std::tuple get_spmd_info() override; + + // Target-specific optimizations + llvm::Value *optimized_reduction(AtomicOpStmt *stmt) override; + + // Kernel argument passing + bool kernel_argument_by_val() const override { + return false; + } + + // Math function emission using FlagOS math library + void emit_extra_unary(UnaryOpStmt *stmt) override; + void visit(BinaryOpStmt *stmt) override; + + // Print support + void create_print(std::string tag, DataType dt, llvm::Value *value) override; + std::tuple create_value_and_type( + llvm::Value *value, + DataType dt); +}; + +/** + * @brief Kernel Code Generation for FlagOS + */ +class KernelCodeGenFlagOS : public KernelCodeGen { + public: + KernelCodeGenFlagOS(const CompileConfig &compile_config, + const Kernel *kernel, + IRNode *ir = nullptr, + TaichiLLVMContext &tlctx) + : KernelCodeGen(compile_config, kernel, ir, tlctx) { + } + + LLVMCompiledTask compile_task( + int task_codegen_id, + const CompileConfig &config, + std::unique_ptr &&module = nullptr, + IRNode *block = nullptr) override; +}; + +} // namespace lang + +} // namespace taichi diff --git a/taichi/inc/archs.inc.h b/taichi/inc/archs.inc.h index d748a7c4eadda..68559b72df5cd 100644 --- a/taichi/inc/archs.inc.h +++ b/taichi/inc/archs.inc.h @@ -15,3 +15,4 @@ PER_ARCH(opencl) // OpenCL, N/A PER_ARCH(amdgpu) // AMD GPU PER_ARCH(vulkan) // Vulkan PER_ARCH(gles) // OpenGL ES +PER_ARCH(flagos) // FlagOS: Unified AI Chip Backend diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h index 87832121f04d6..c727b4a7fc8f5 100644 --- a/taichi/program/compile_config.h +++ b/taichi/program/compile_config.h @@ -102,6 +102,10 @@ struct CompileConfig { size_t cuda_stack_limit{0}; + // FlagOS backend options: + std::string flagos_chip{ + "generic"}; // Target chip: mlu370, ascend910, dcu, etc. + CompileConfig(); void fit(); diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp index 52a6b51db109d..db2d065612b1f 100644 --- a/taichi/program/program.cpp +++ b/taichi/program/program.cpp @@ -42,6 +42,10 @@ #include "taichi/rhi/metal/metal_api.h" #endif // TI_WITH_METAL +#ifdef TI_WITH_FLAGOS +#include "taichi/runtime/program_impls/flagos/flagos_program.h" +#endif // TI_WITH_FLAGOS + #if defined(_M_X64) || defined(__x86_64) // For _MM_SET_FLUSH_ZERO_MODE #include @@ -79,7 +83,14 @@ Program::Program(Arch desired_arch) : snode_rw_accessors_bank_(this) { profiler = make_profiler(config.arch, config.kernel_profiler); if (arch_uses_llvm(config.arch)) { #ifdef TI_WITH_LLVM - if (config.arch != Arch::dx12) { + if (config.arch == Arch::flagos) { +#ifdef TI_WITH_FLAGOS + program_impl_ = + std::make_unique(config, profiler.get()); +#else + TI_ERROR("This taichi is not compiled with FlagOS"); +#endif + } else if (config.arch != Arch::dx12) { program_impl_ = std::make_unique(config, profiler.get()); } else { // NOTE: use Dx12ProgramImpl to avoid using LlvmRuntimeExecutor for dx12. diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp index 4715b51f02914..b22d9dd11112d 100644 --- a/taichi/python/export_lang.cpp +++ b/taichi/python/export_lang.cpp @@ -256,7 +256,8 @@ void export_lang(py::module &m) { &CompileConfig::offline_cache_cleaning_factor) .def_readwrite("num_compile_threads", &CompileConfig::num_compile_threads) .def_readwrite("vk_api_version", &CompileConfig::vk_api_version) - .def_readwrite("cuda_stack_limit", &CompileConfig::cuda_stack_limit); + .def_readwrite("cuda_stack_limit", &CompileConfig::cuda_stack_limit) + .def_readwrite("flagos_chip", &CompileConfig::flagos_chip); m.def("reset_default_compile_config", [&]() { default_compile_config = CompileConfig(); }); diff --git a/taichi/rhi/CMakeLists.txt b/taichi/rhi/CMakeLists.txt index 4b88b7d974c0d..3e0c0b3b7abb1 100644 --- a/taichi/rhi/CMakeLists.txt +++ b/taichi/rhi/CMakeLists.txt @@ -81,6 +81,10 @@ if (TI_WITH_VULKAN) target_link_libraries(${TAICHI_DEVICE_API} PUBLIC vulkan_rhi) endif () +if(TI_WITH_FLAGOS) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_WITH_FLAGOS") +endif() + if(TI_WITH_LLVM) add_subdirectory(cpu) target_link_libraries(${TAICHI_DEVICE_API} PUBLIC cpu_rhi) @@ -100,6 +104,11 @@ if(TI_WITH_LLVM) target_link_libraries(${TAICHI_DEVICE_API} PUBLIC dx12_rhi) endif() + if (TI_WITH_FLAGOS) + add_subdirectory(flagos) + target_link_libraries(${TAICHI_DEVICE_API} PUBLIC flagos_rhi) + endif() + add_subdirectory(llvm) target_link_libraries(${TAICHI_DEVICE_API} PUBLIC llvm_rhi) endif() diff --git a/taichi/rhi/arch.cpp b/taichi/rhi/arch.cpp index 7e7bb72a5d2b0..dc2ae7d1c4f61 100644 --- a/taichi/rhi/arch.cpp +++ b/taichi/rhi/arch.cpp @@ -53,7 +53,7 @@ bool arch_is_cuda(Arch arch) { bool arch_uses_llvm(Arch arch) { return (arch == Arch::x64 || arch == Arch::arm64 || arch == Arch::cuda || - arch == Arch::dx12 || arch == Arch::amdgpu); + arch == Arch::dx12 || arch == Arch::amdgpu || arch == Arch::flagos); } bool arch_is_gpu(Arch arch) { @@ -86,6 +86,9 @@ int default_simd_width(Arch arch) { return 32; } else if (arch == Arch::arm64) { return 4; + } else if (arch == Arch::flagos) { + // FlagOS default SIMD width, can be configured per chip + return 32; } else { RHI_NOT_IMPLEMENTED; return -1; diff --git a/taichi/rhi/flagos/CMakeLists.txt b/taichi/rhi/flagos/CMakeLists.txt new file mode 100644 index 0000000000000..86796ff6748b3 --- /dev/null +++ b/taichi/rhi/flagos/CMakeLists.txt @@ -0,0 +1,30 @@ +# ./taichi/rhi/flagos/CMakeLists.txt + +set(FLAGOS_RHI flagos_rhi) +add_library(${FLAGOS_RHI}) + +target_sources(${FLAGOS_RHI} + PRIVATE + flagos_device.cpp + ) + +target_include_directories(${FLAGOS_RHI} + PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/external/eigen + ${PROJECT_SOURCE_DIR}/external/spdlog/include + ) + +# FlagOS SDK integration +if(TI_WITH_FLAGOS) + target_compile_definitions(${FLAGOS_RHI} PRIVATE TI_WITH_FLAGOS) + + # Find FlagOS SDK + find_package(FlagOS QUIET) + if(FlagOS_FOUND) + target_link_libraries(${FLAGOS_RHI} PRIVATE FlagOS::FlagTree) + target_include_directories(${FLAGOS_RHI} PRIVATE ${FlagOS_INCLUDE_DIRS}) + else() + message(WARNING "FlagOS SDK not found, building with stub implementation") + endif() +endif() diff --git a/taichi/rhi/flagos/flagos_device.cpp b/taichi/rhi/flagos/flagos_device.cpp new file mode 100644 index 0000000000000..b0e255a535e6f --- /dev/null +++ b/taichi/rhi/flagos/flagos_device.cpp @@ -0,0 +1,349 @@ +#include "taichi/rhi/flagos/flagos_device.h" +#include "taichi/rhi/llvm/device_memory_pool.h" + +#include "taichi/jit/jit_module.h" +#include "taichi/common/logging.h" + +namespace taichi { +namespace lang { + +namespace flagos { + +// FlagOS context wrapper implementation +struct FlagosDevice::FlagosContextWrapper { + // Placeholder for actual FlagOS context + // This will be replaced with actual FlagOS SDK calls + bool initialized{false}; + + void initialize(const std::string &chip_name) { + TI_INFO("Initializing FlagOS backend with chip: {}", chip_name); + // TODO: Call FlagOS SDK initialization + // flagos_init(chip_name.c_str()); + initialized = true; + } + + void *allocate_memory(size_t size, bool managed) { + // TODO: Use FlagOS memory allocation + // return flagos_malloc(size, managed); + TI_NOT_IMPLEMENTED; + return nullptr; + } + + void free_memory(void *ptr) { + // TODO: Use FlagOS memory deallocation + // flagos_free(ptr); + } + + void memcpy_host_to_device(void *dst, const void *src, size_t size) { + // TODO: Use FlagOS H2D copy + TI_NOT_IMPLEMENTED; + } + + void memcpy_device_to_host(void *dst, const void *src, size_t size) { + // TODO: Use FlagOS D2H copy + TI_NOT_IMPLEMENTED; + } + + void memcpy_device_to_device(void *dst, const void *src, size_t size) { + // TODO: Use FlagOS D2D copy + TI_NOT_IMPLEMENTED; + } + + void memset(void *ptr, int value, size_t size) { + // TODO: Use FlagOS memset + TI_NOT_IMPLEMENTED; + } + + void synchronize() { + // TODO: Use FlagOS synchronization + TI_NOT_IMPLEMENTED; + } + + size_t get_total_memory() { + // TODO: Query FlagOS device memory + return 0; + } + + void launch_kernel(const std::string &kernel_name, + void **args, + uint32_t grid_dim, + uint32_t block_dim) { + // TODO: Launch kernel via FlagOS + TI_NOT_IMPLEMENTED; + } +}; + +FlagosDevice::FlagosDevice() + : context_(std::make_unique()) { + // Initialize the device memory pool + DeviceMemoryPool::get_instance(false /*merge_upon_release*/); + + // Get target chip from environment variable or use default + const char *chip_env = std::getenv("TI_FLAGOS_CHIP"); + if (chip_env) { + target_chip_ = chip_env; + } + + TI_INFO("FlagOS Device created, target chip: {}", target_chip_); +} + +FlagosDevice::~FlagosDevice() { + clear(); +} + +void FlagosDevice::initialize(const std::string &chip_name) { + target_chip_ = chip_name; + context_->initialize(chip_name); +} + +FlagosDevice::AllocInfo FlagosDevice::get_alloc_info( + const DeviceAllocation handle) { + validate_device_alloc(handle); + return allocations_[handle.alloc_id]; +} + +RhiResult FlagosDevice::allocate_memory(const AllocParams ¶ms, + DeviceAllocation *out_devalloc) { + AllocInfo info; + + auto &mem_pool = DeviceMemoryPool::get_instance(); + + bool managed = params.host_read || params.host_write; + + // Try to use FlagOS allocation first, fallback to memory pool + void *ptr = nullptr; + if (context_->initialized) { + ptr = context_->allocate_memory(params.size, managed); + } + + if (ptr == nullptr) { + ptr = mem_pool.allocate(params.size, DeviceMemoryPool::page_size, managed); + } + + if (ptr == nullptr) { + return RhiResult::out_of_memory; + } + + info.ptr = ptr; + info.size = params.size; + info.is_imported = false; + info.use_cached = false; + info.use_preallocated = false; + + if (info.ptr == nullptr) { + return RhiResult::out_of_memory; + } + + // Initialize memory to zero + if (context_->initialized) { + context_->memset(info.ptr, 0, info.size); + } else { + std::memset(info.ptr, 0, info.size); + } + + *out_devalloc = DeviceAllocation{}; + out_devalloc->alloc_id = allocations_.size(); + out_devalloc->device = this; + + allocations_.push_back(info); + return RhiResult::success; +} + +DeviceAllocation FlagosDevice::allocate_memory_runtime( + const LlvmRuntimeAllocParams ¶ms) { + AllocInfo info; + info.size = taichi::iroundup(params.size, taichi_page_size); + if (params.host_read || params.host_write) { + TI_NOT_IMPLEMENTED + } else { + info.ptr = + DeviceMemoryPool::get_instance().allocate_with_cache(this, params); + TI_ASSERT(info.ptr != nullptr); + + if (context_->initialized) { + context_->memset(info.ptr, 0, info.size); + } + } + info.is_imported = false; + info.use_cached = true; + info.use_preallocated = true; + + DeviceAllocation alloc; + alloc.alloc_id = allocations_.size(); + alloc.device = this; + + allocations_.push_back(info); + return alloc; +} + +uint64_t *FlagosDevice::allocate_llvm_runtime_memory_jit( + const LlvmRuntimeAllocParams ¶ms) { + params.runtime_jit->call( + "runtime_memory_allocate_aligned", params.runtime, params.size, + taichi_page_size, params.result_buffer); + + wait_idle(); + + uint64 *ret{nullptr}; + // Read back the result + if (context_->initialized) { + context_->memcpy_device_to_host(&ret, params.result_buffer, sizeof(uint64)); + } + return ret; +} + +void FlagosDevice::dealloc_memory(DeviceAllocation handle) { + // After reset, all allocations are invalid + if (allocations_.empty()) { + return; + } + + validate_device_alloc(handle); + AllocInfo &info = allocations_[handle.alloc_id]; + if (info.ptr == nullptr) { + TI_ERROR("the DeviceAllocation is already deallocated"); + } + TI_ASSERT(!info.is_imported); + if (info.use_cached) { + DeviceMemoryPool::get_instance().release(info.size, (uint64_t *)info.ptr, + false); + } else if (!info.use_preallocated) { + if (context_->initialized) { + context_->free_memory(info.ptr); + } else { + DeviceMemoryPool::get_instance().release(info.size, info.ptr); + } + } + info.ptr = nullptr; +} + +RhiResult FlagosDevice::map(DeviceAllocation alloc, void **mapped_ptr) { + AllocInfo &info = allocations_[alloc.alloc_id]; + size_t size = info.size; + info.mapped = new char[size]; + + if (context_->initialized) { + context_->memcpy_device_to_host(info.mapped, info.ptr, size); + } + + *mapped_ptr = info.mapped; + return RhiResult::success; +} + +void FlagosDevice::unmap(DeviceAllocation alloc) { + AllocInfo &info = allocations_[alloc.alloc_id]; + + if (context_->initialized) { + context_->memcpy_host_to_device(info.ptr, info.mapped, info.size); + } + + delete[] static_cast(info.mapped); + info.mapped = nullptr; +} + +void FlagosDevice::memcpy_internal(DevicePtr dst, + DevicePtr src, + uint64_t size) { + void *dst_ptr = + static_cast(allocations_[dst.alloc_id].ptr) + dst.offset; + void *src_ptr = + static_cast(allocations_[src.alloc_id].ptr) + src.offset; + + if (context_->initialized) { + context_->memcpy_device_to_device(dst_ptr, src_ptr, size); + } +} + +DeviceAllocation FlagosDevice::import_memory(void *ptr, size_t size) { + AllocInfo info; + info.ptr = ptr; + info.size = size; + info.is_imported = true; + + DeviceAllocation alloc; + alloc.alloc_id = allocations_.size(); + alloc.device = this; + + allocations_.push_back(info); + return alloc; +} + +std::size_t FlagosDevice::get_total_memory() { + if (context_->initialized) { + return context_->get_total_memory(); + } + return 0; +} + +void FlagosDevice::wait_idle() { + if (context_->initialized) { + context_->synchronize(); + } +} + +void FlagosDevice::clear() { + allocations_.clear(); + context_->initialized = false; +} + +RhiResult FlagosDevice::upload_data(DevicePtr *device_ptr, + const void **data, + size_t *size, + int num_alloc) noexcept { + for (int i = 0; i < num_alloc; i++) { + void *dst_ptr = + static_cast(allocations_[device_ptr[i].alloc_id].ptr) + + device_ptr[i].offset; + if (context_->initialized) { + context_->memcpy_host_to_device(dst_ptr, data[i], size[i]); + } + } + return RhiResult::success; +} + +RhiResult FlagosDevice::readback_data( + DevicePtr *device_ptr, + void **data, + size_t *size, + int num_alloc, + const std::vector &wait_sema) noexcept { + for (int i = 0; i < num_alloc; i++) { + void *src_ptr = + static_cast(allocations_[device_ptr[i].alloc_id].ptr) + + device_ptr[i].offset; + if (context_->initialized) { + context_->memcpy_device_to_host(data[i], src_ptr, size[i]); + } + } + return RhiResult::success; +} + +void FlagosDevice::launch_kernel(const std::string &kernel_name, + void **args, + uint32_t grid_dim, + uint32_t block_dim) { + if (context_->initialized) { + context_->launch_kernel(kernel_name, args, grid_dim, block_dim); + } else { + TI_ERROR("FlagOS context not initialized"); + } +} + +bool FlagosDevice::is_chip_supported(const std::string &chip_name) { + // TODO: Query FlagOS for supported chips + static const std::set supported_chips = { + "mlu370", "mlu590", "ascend910", "ascend310", "dcu", "gcu", "generic"}; + return supported_chips.count(chip_name) > 0; +} + +std::vector FlagosDevice::get_supported_chips() { + // TODO: Query FlagOS for available chips + return {"mlu370", "mlu590", "ascend910", "ascend310", + "dcu", "gcu", "generic"}; +} + +} // namespace flagos + +} // namespace lang + +} // namespace taichi diff --git a/taichi/rhi/flagos/flagos_device.h b/taichi/rhi/flagos/flagos_device.h new file mode 100644 index 0000000000000..0b212c58d758c --- /dev/null +++ b/taichi/rhi/flagos/flagos_device.h @@ -0,0 +1,222 @@ +#pragma once + +#include +#include +#include +#include + +#include "taichi/common/core.h" +#include "taichi/rhi/llvm/llvm_device.h" + +namespace taichi { +namespace lang { + +namespace flagos { + +/** + * @brief FlagOS Device abstraction for Taichi + * + * FlagOS is a unified AI system software stack that provides a common interface + * for various AI chips. This device implementation bridges Taichi's LLVM + * backend with FlagOS's FlagTree compiler. + */ + +class FlagosPipeline : public Pipeline { + public: + ~FlagosPipeline() override { + } + + // Compiled kernel handle from FlagTree + void *kernel_handle_{nullptr}; + std::string kernel_name_; +}; + +class FlagosCommandList : public CommandList { + public: + ~FlagosCommandList() override { + } + + void bind_pipeline(Pipeline *p) noexcept override { + TI_NOT_IMPLEMENTED; + } + RhiResult bind_shader_resources(ShaderResourceSet *res, + int set_index = 0) noexcept override{ + TI_NOT_IMPLEMENTED} RhiResult + bind_raster_resources(RasterResources *res) noexcept override { + TI_NOT_IMPLEMENTED + } + void buffer_barrier(DevicePtr ptr, size_t size) noexcept override { + TI_NOT_IMPLEMENTED + } + void buffer_barrier(DeviceAllocation alloc) noexcept override { + TI_NOT_IMPLEMENTED + } + void memory_barrier() noexcept override { + TI_NOT_IMPLEMENTED; + } + void buffer_copy(DevicePtr dst, + DevicePtr src, + size_t size) noexcept override { + TI_NOT_IMPLEMENTED + } + void buffer_fill(DevicePtr ptr, size_t size, uint32_t data) noexcept override{ + TI_NOT_IMPLEMENTED} RhiResult + dispatch(uint32_t x, uint32_t y = 1, uint32_t z = 1) noexcept override { + TI_NOT_IMPLEMENTED + } +}; + +class FlagosStream : public Stream { + public: + ~FlagosStream() override { + } + + RhiResult new_command_list(CommandList **out_cmdlist) noexcept override{ + TI_NOT_IMPLEMENTED} StreamSemaphore + submit(CommandList *cmdlist, + const std::vector &wait_semaphores = {}) override{ + TI_NOT_IMPLEMENTED} StreamSemaphore + submit_synced( + CommandList *cmdlist, + const std::vector &wait_semaphores = {}) override { + TI_NOT_IMPLEMENTED + } + + void command_sync() override { + TI_NOT_IMPLEMENTED; + } +}; + +class FlagosDevice : public LlvmDevice { + public: + struct AllocInfo { + void *ptr{nullptr}; + size_t size{0}; + bool is_imported{false}; + bool use_preallocated{true}; + bool use_cached{false}; + void *mapped{nullptr}; + }; + + AllocInfo get_alloc_info(const DeviceAllocation handle); + + FlagosDevice(); + ~FlagosDevice() override; + + // Device capabilities + Arch arch() const override { + return Arch::flagos; + } + + // Memory management + RhiResult allocate_memory(const AllocParams ¶ms, + DeviceAllocation *out_devalloc) override; + DeviceAllocation allocate_memory_runtime( + const LlvmRuntimeAllocParams ¶ms) override; + void dealloc_memory(DeviceAllocation handle) override; + + uint64_t *allocate_llvm_runtime_memory_jit( + const LlvmRuntimeAllocParams ¶ms) override; + + RhiResult upload_data(DevicePtr *device_ptr, + const void **data, + size_t *size, + int num_alloc = 1) noexcept override; + + RhiResult readback_data( + DevicePtr *device_ptr, + void **data, + size_t *size, + int num_alloc = 1, + const std::vector &wait_sema = {}) noexcept override; + + ShaderResourceSet *create_resource_set() final { + TI_NOT_IMPLEMENTED; + } + + RhiResult create_pipeline(Pipeline **out_pipeline, + const PipelineSourceDesc &src, + std::string name, + PipelineCache *cache) noexcept final { + TI_NOT_IMPLEMENTED; + } + + RhiResult map_range(DevicePtr ptr, uint64_t size, void **mapped_ptr) final; + RhiResult map(DeviceAllocation alloc, void **mapped_ptr) final; + + void unmap(DevicePtr ptr) override { + TI_NOT_IMPLEMENTED; + } + void unmap(DeviceAllocation alloc) override; + + void memcpy_internal(DevicePtr dst, DevicePtr src, uint64_t size) override; + + DeviceAllocation import_memory(void *ptr, size_t size) override; + + // LLVM Device interface + void *get_memory_addr(DeviceAllocation devalloc) override { + return get_alloc_info(devalloc).ptr; + } + + std::size_t get_total_memory() override; + + Stream *get_compute_stream() override { + TI_NOT_IMPLEMENTED; + } + + void wait_idle() override; + + void clear() override; + + // FlagOS specific methods + /** + * @brief Initialize FlagOS backend with specific chip target + * @param chip_name Target chip name (e.g., "mlu370", "ascend910", "dcu") + */ + void initialize(const std::string &chip_name); + + /** + * @brief Get the target chip name + */ + std::string get_target_chip() const { + return target_chip_; + } + + /** + * @brief Check if a specific chip is supported + */ + static bool is_chip_supported(const std::string &chip_name); + + /** + * @brief Get list of supported chips + */ + static std::vector get_supported_chips(); + + /** + * @brief Launch a compiled kernel + */ + void launch_kernel(const std::string &kernel_name, + void **args, + uint32_t grid_dim, + uint32_t block_dim); + + private: + std::vector allocations_; + std::string target_chip_{"generic"}; + + // FlagOS context + struct FlagosContextWrapper; + std::unique_ptr context_; + + void validate_device_alloc(const DeviceAllocation alloc) { + if (allocations_.size() <= alloc.alloc_id) { + TI_ERROR("invalid DeviceAllocation"); + } + } +}; + +} // namespace flagos + +} // namespace lang + +} // namespace taichi diff --git a/taichi/runtime/program_impls/flagos/CMakeLists.txt b/taichi/runtime/program_impls/flagos/CMakeLists.txt new file mode 100644 index 0000000000000..001926a5491b5 --- /dev/null +++ b/taichi/runtime/program_impls/flagos/CMakeLists.txt @@ -0,0 +1,31 @@ +# ./taichi/runtime/program_impls/flagos/CMakeLists.txt + +set(FLAGOS_PROGRAM flagos_program) +add_library(${FLAGOS_PROGRAM}) + +target_sources(${FLAGOS_PROGRAM} + PRIVATE + flagos_program.cpp + flagos_kernel_compiler.cpp + flagos_kernel_launcher.cpp + ) + +target_include_directories(${FLAGOS_PROGRAM} + PRIVATE + ${PROJECT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/external/spdlog/include + ${LLVM_INCLUDE_DIRS} + ) + +target_link_libraries(${FLAGOS_PROGRAM} + PRIVATE + taichi_util + taichi_ir + llvm_program_impl + flagos_rhi + flagos_codegen + ) + +if(TI_WITH_FLAGOS) + target_compile_definitions(${FLAGOS_PROGRAM} PRIVATE TI_WITH_FLAGOS) +endif() diff --git a/taichi/runtime/program_impls/flagos/flagos_kernel_compiler.cpp b/taichi/runtime/program_impls/flagos/flagos_kernel_compiler.cpp new file mode 100644 index 0000000000000..e29a1769b5304 --- /dev/null +++ b/taichi/runtime/program_impls/flagos/flagos_kernel_compiler.cpp @@ -0,0 +1,68 @@ +#include "taichi/runtime/program_impls/flagos/flagos_kernel_compiler.h" + +#include "taichi/program/kernel.h" +#include "taichi/analysis/offline_cache_util.h" +#include "taichi/runtime/llvm/llvm_context.h" + +namespace taichi { +namespace lang { +namespace flagos { + +KernelCompiler::KernelCompiler(Config config) : config_(std::move(config)) { +} + +LLVM::CompiledKernelData KernelCompiler::compile( + const CompileConfig &compile_config, + const DeviceCapabilityConfig &device_caps, + const Kernel &kernel, + IRNode *ir) { + return compile_with_flagtree(compile_config, kernel, ir); +} + +LLVM::CompiledKernelData KernelCompiler::compile( + const CompileConfig &compile_config, + const DeviceCapabilityConfig &device_caps, + const Kernel &kernel, + IRNode *ir, + const LLVM::CompiledKernelData &cached_data) { + // Check if cached data is compatible + if (cached_data.arch() == compile_config.arch) { + TI_TRACE("Using cached FlagOS kernel: {}", kernel.get_name()); + return cached_data; + } + + return compile_with_flagtree(compile_config, kernel, ir); +} + +LLVM::CompiledKernelData KernelCompiler::compile_with_flagtree( + const CompileConfig &compile_config, + const Kernel &kernel, + IRNode *ir) { + TI_TRACE("Compiling FlagOS kernel: {}", kernel.get_name()); + + // Create FlagOS code generator + auto codegen = + KernelCodeGenFlagOS(compile_config, &kernel, ir, *config_.tlctx); + + // Compile to LLVM IR + auto compiled_tasks = codegen.run_compilation(); + + // Create compiled kernel data + LLVM::CompiledKernelData::InternalData data; + data.arch = compile_config.arch; + data.kernel_key = get_cache_key(&kernel); + data.tasks = std::move(compiled_tasks.tasks); + data.args = std::move(compiled_tasks.args); + data.module = std::move(compiled_tasks.module); + + // TODO: Use FlagTree to compile LLVM IR to target chip code + // auto flagtree = FlagTreeCompiler(); + // auto chip_code = flagtree.compile(data.module, compile_config.flagos_chip); + // data.chip_binary = std::move(chip_code); + + return LLVM::CompiledKernelData(std::move(data)); +} + +} // namespace flagos +} // namespace lang +} // namespace taichi diff --git a/taichi/runtime/program_impls/flagos/flagos_kernel_compiler.h b/taichi/runtime/program_impls/flagos/flagos_kernel_compiler.h new file mode 100644 index 0000000000000..484382d7df1af --- /dev/null +++ b/taichi/runtime/program_impls/flagos/flagos_kernel_compiler.h @@ -0,0 +1,56 @@ +#pragma once + +#include "taichi/codegen/llvm/kernel_compiler.h" +#include "taichi/codegen/flagos/codegen_flagos.h" + +namespace taichi { +namespace lang { +namespace flagos { + +/** + * @brief FlagOS Kernel Compiler + * + * Compiles Taichi kernels to LLVM IR for FlagOS backend. + * Uses FlagTree to compile LLVM IR to target chip code. + */ + +class KernelCompiler : public LLVM::KernelCompiler { + public: + using IRNode = taichi::lang::IRNode; + using Config = LLVM::KernelCompiler::Config; + + explicit KernelCompiler(Config config); + + /** + * @brief Compile a kernel to FlagOS executable format + */ + LLVM::CompiledKernelData compile(const CompileConfig &compile_config, + const DeviceCapabilityConfig &device_caps, + const Kernel &kernel, + IRNode *ir) override; + + /** + * @brief Compile a kernel with cached data + */ + LLVM::CompiledKernelData compile( + const CompileConfig &compile_config, + const DeviceCapabilityConfig &device_caps, + const Kernel &kernel, + IRNode *ir, + const LLVM::CompiledKernelData &cached_data) override; + + private: + Config config_; + + /** + * @brief Compile kernel using FlagTree + */ + LLVM::CompiledKernelData compile_with_flagtree( + const CompileConfig &compile_config, + const Kernel &kernel, + IRNode *ir); +}; + +} // namespace flagos +} // namespace lang +} // namespace taichi diff --git a/taichi/runtime/program_impls/flagos/flagos_kernel_launcher.cpp b/taichi/runtime/program_impls/flagos/flagos_kernel_launcher.cpp new file mode 100644 index 0000000000000..0230594b930c4 --- /dev/null +++ b/taichi/runtime/program_impls/flagos/flagos_kernel_launcher.cpp @@ -0,0 +1,59 @@ +#include "taichi/runtime/program_impls/flagos/flagos_kernel_launcher.h" + +#include "taichi/rhi/flagos/flagos_device.h" +#include "taichi/runtime/llvm/llvm_runtime_executor.h" +#include "taichi/program/kernel.h" + +namespace taichi { +namespace lang { +namespace flagos { + +KernelLauncher::KernelLauncher(Config config) : config_(std::move(config)) { +} + +void KernelLauncher::launch_kernel( + LlvmLaunchArgBuilder &arg_builder, + Kernel *kernel, + const LLVM::CompiledKernelData &compiled_kernel) { + // Get the FlagOS device + auto *device = + static_cast(config_.executor->get_compute_device()); + TI_ASSERT(device); + + // Get kernel data + const auto &kernel_data = compiled_kernel.get_internal_data(); + const auto &task_funcs = kernel_data.tasks; + const auto &args = kernel_data.args; + + // Prepare launch arguments + std::vector arg_pointers; + arg_pointers.reserve(args.size()); + + for (const auto &arg : args) { + arg_pointers.push_back(const_cast(arg.data)); + } + + // Launch each task (offloaded kernel) + for (const auto &task : task_funcs) { + const std::string &kernel_name = task.name; + uint32_t grid_dim = task.grid_dim; + uint32_t block_dim = task.block_dim; + + TI_TRACE("Launching FlagOS kernel: {}, grid_dim: {}, block_dim: {}", + kernel_name, grid_dim, block_dim); + + // Launch kernel via FlagOS device + // This will call FlagTree to compile and execute the kernel + device->launch_kernel(kernel_name, arg_pointers.data(), grid_dim, + block_dim); + } + + // Synchronize if needed + if (kernel->program->compile_config().sync_after_kernel_launch) { + device->wait_idle(); + } +} + +} // namespace flagos +} // namespace lang +} // namespace taichi diff --git a/taichi/runtime/program_impls/flagos/flagos_kernel_launcher.h b/taichi/runtime/program_impls/flagos/flagos_kernel_launcher.h new file mode 100644 index 0000000000000..58c0913648514 --- /dev/null +++ b/taichi/runtime/program_impls/flagos/flagos_kernel_launcher.h @@ -0,0 +1,35 @@ +#pragma once + +#include "taichi/runtime/llvm/kernel_launcher.h" + +namespace taichi { +namespace lang { +namespace flagos { + +/** + * @brief FlagOS Kernel Launcher + * + * This class handles kernel launching for FlagOS-supported AI chips. + * It integrates with the FlagOS runtime to execute compiled kernels. + */ + +class KernelLauncher : public LLVM::KernelLauncher { + public: + using Config = LLVM::KernelLauncher::Config; + + explicit KernelLauncher(Config config); + + /** + * @brief Launch a kernel on FlagOS device + */ + void launch_kernel(LlvmLaunchArgBuilder &arg_builder, + Kernel *kernel, + const LLVM::CompiledKernelData &compiled_kernel) override; + + private: + Config config_; +}; + +} // namespace flagos +} // namespace lang +} // namespace taichi diff --git a/taichi/runtime/program_impls/flagos/flagos_program.cpp b/taichi/runtime/program_impls/flagos/flagos_program.cpp new file mode 100644 index 0000000000000..488f0d85626a0 --- /dev/null +++ b/taichi/runtime/program_impls/flagos/flagos_program.cpp @@ -0,0 +1,111 @@ +#include "taichi/runtime/program_impls/flagos/flagos_program.h" + +#include "taichi/codegen/flagos/codegen_flagos.h" +#include "taichi/runtime/llvm/llvm_runtime_executor.h" +#include "taichi/codegen/llvm/struct_llvm.h" +#include "taichi/program/program.h" +#include "taichi/runtime/llvm/llvm_aot_module_builder.h" +#include "taichi/runtime/program_impls/flagos/flagos_kernel_compiler.h" +#include "taichi/runtime/program_impls/flagos/flagos_kernel_launcher.h" + +namespace taichi { +namespace lang { + +FlagosProgramImpl::FlagosProgramImpl(CompileConfig &config, + KernelProfilerBase *profiler) + : LlvmProgramImpl(config, profiler) { + // Initialize FlagOS-specific settings + if (config.flagos_chip.empty()) { + // Try to get from environment variable + const char *chip_env = std::getenv("TI_FLAGOS_CHIP"); + if (chip_env) { + config.flagos_chip = chip_env; + } else { + config.flagos_chip = "generic"; // Default chip + } + } + + TI_INFO("FlagosProgramImpl created with chip: {}", config.flagos_chip); + + // Validate chip name + if (!flagos::FlagosDevice::is_chip_supported(config.flagos_chip)) { + TI_WARN("FlagOS chip '{}' may not be fully supported", config.flagos_chip); + } +} + +FlagosProgramImpl::~FlagosProgramImpl() { + // Cleanup is handled by base class +} + +void FlagosProgramImpl::initialize_flagos_backend() { + auto *device = static_cast(llvm_device()); + if (device) { + device->initialize(config_.flagos_chip); + TI_INFO("FlagOS backend initialized for chip: {}", config_.flagos_chip); + } +} + +bool FlagosProgramImpl::is_initialized() const { + auto *device = static_cast(llvm_device()); + if (device) { + return flagos::FlagosDevice::is_chip_supported(device->get_target_chip()); + } + return false; +} + +void FlagosProgramImpl::compile_snode_tree_types(SNodeTree *tree) { + // Delegate to base class implementation + LlvmProgramImpl::compile_snode_tree_types(tree); +} + +void FlagosProgramImpl::materialize_snode_tree(SNodeTree *tree, + uint64 *result_buffer) { + // Initialize FlagOS backend on first materialization + static bool flagos_initialized = false; + if (!flagos_initialized) { + initialize_flagos_backend(); + flagos_initialized = true; + } + + // Delegate to base class implementation + LlvmProgramImpl::materialize_snode_tree(tree, result_buffer); +} + +void FlagosProgramImpl::materialize_runtime(KernelProfilerBase *profiler, + uint64 **result_buffer_ptr) { + // First initialize base LLVM runtime + LlvmProgramImpl::materialize_runtime(profiler, result_buffer_ptr); + + // Then initialize FlagOS-specific context + initialize_flagos_backend(); +} + +std::unique_ptr FlagosProgramImpl::make_kernel_compiler() { + flagos::KernelCompiler::Config cfg; + cfg.tlctx = get_llvm_context(); + return std::make_unique(std::move(cfg)); +} + +std::unique_ptr FlagosProgramImpl::make_kernel_launcher() { + flagos::KernelLauncher::Config cfg; + cfg.executor = get_runtime_executor(); + return std::make_unique(std::move(cfg)); +} + +std::unique_ptr FlagosProgramImpl::make_aot_module_builder( + const DeviceCapabilityConfig &caps) { + // Use LLVM AOT module builder for now + // TODO: Create FlagOS-specific AOT module builder + return std::make_unique( + get_kernel_compilation_manager(), config_, this); +} + +FlagosProgramImpl *get_flagos_program(Program *prog) { + auto *result = dynamic_cast(prog->get_program_impl()); + TI_ASSERT(result); + return result; +} + +} // namespace lang + +} // namespace taichi diff --git a/taichi/runtime/program_impls/flagos/flagos_program.h b/taichi/runtime/program_impls/flagos/flagos_program.h new file mode 100644 index 0000000000000..68e35e1635335 --- /dev/null +++ b/taichi/runtime/program_impls/flagos/flagos_program.h @@ -0,0 +1,101 @@ +#pragma once + +#include +#include +#include + +#include "taichi/runtime/program_impls/llvm/llvm_program.h" +#include "taichi/rhi/flagos/flagos_device.h" + +namespace taichi { +namespace lang { + +// Forward declarations +class KernelCompiler; +class KernelLauncher; + +/** + * @brief FlagOS Program Implementation + * + * This class extends LlvmProgramImpl to provide program execution + * support for the FlagOS unified AI chip backend. + * + * Features: + * - Multi-chip support (MLU, Ascend, DCU, GCU, etc.) + * - Integration with FlagTree compiler + * - Optimized kernel compilation and launching + */ + +class FlagosProgramImpl : public LlvmProgramImpl { + public: + FlagosProgramImpl(CompileConfig &config, KernelProfilerBase *profiler); + + /* ------------------------------------ */ + /* ---- Compilation Interfaces ---- */ + /* ------------------------------------ */ + + void compile_snode_tree_types(SNodeTree *tree) override; + void materialize_snode_tree(SNodeTree *tree, uint64 *result_buffer) override; + + /* -------------------------------- */ + /* ---- Runtime Interfaces ---- */ + /* -------------------------------- */ + + void materialize_runtime(KernelProfilerBase *profiler, + uint64 **result_buffer_ptr) override; + + void finalize() override { + LlvmProgramImpl::finalize(); + } + + void synchronize() override { + LlvmProgramImpl::synchronize(); + } + + Device *get_compute_device() override { + return LlvmProgramImpl::get_compute_device(); + } + + /* -------------------------------- */ + /* ---- FlagOS Specific ---- */ + /* -------------------------------- */ + + /** + * @brief Get the target chip for this program + */ + std::string get_target_chip() const { + return config_.flagos_chip; + } + + /** + * @brief Set the target chip for this program + */ + void set_target_chip(const std::string &chip) { + config_.flagos_chip = chip; + } + + /** + * @brief Check if FlagOS backend is properly initialized + */ + bool is_initialized() const; + + ~FlagosProgramImpl() override; + + protected: + std::unique_ptr make_kernel_compiler() override; + std::unique_ptr make_kernel_launcher() override; + std::unique_ptr make_aot_module_builder( + const DeviceCapabilityConfig &caps) override; + + private: + void initialize_flagos_backend(); +}; + +/** + * @brief Get FlagOS program from Taichi Program + */ +FlagosProgramImpl *get_flagos_program(Program *prog); + +} // namespace lang + +} // namespace taichi