| [mlir][ptr] Extend ptr_add operation to support shaped operands (#156374) This patch extends ptr_add to work with shaped types with value semantics, both for the offsets and base. Concretely this patch makes the following changes: - Supports scalar-to-scalar, scalar-to-shaped, shaped-to-scalar, and shaped-to-shaped combinations - Adds InferTypeOpInterface for automatic result type deduction - Adds tests for LLVM IR translation with vector operands Example: mlir func.func @ptr_add_tensor_2d(%ptrs: tensor<4x8x!ptr.ptr<#ptr.generic_space>>, %offsets: tensor<4x8xindex>) -> tensor<4x8x!ptr.ptr<#ptr.generic_space>> { %res = ptr.ptr_add %ptrs, %offsets : tensor<4x8x!ptr.ptr<#ptr.generic_space>>, tensor<4x8xindex> %res1 = ptr.ptr_add nuw %ptrs, %offsets : tensor<4x8x!ptr.ptr<#ptr.generic_space>>, tensor<4x8xindex> return %res : tensor<4x8x!ptr.ptr<#ptr.generic_space>> } The motivation behind this patch is to lay the groundwork for enabling triton styled loads and stores, and their variants. --------- Co-authored-by: Mehdi Amini <joker.eph@gmail.com> | 11 个月前 |