[llvm-branch-commits] [mlir] [mlir][MPIToLLVM] Take the descriptor index type from the descriptor (PR #218838)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue Aug 25 22:31:51 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Christian Ulmann (Dinistro)
<details>
<summary>Changes</summary>
`getRawPtrAndSize` extracted the memref descriptor's offset as `i64` and unconditionally truncated the extents to `i32`. Both assume a 64-bit index: with a 32-bit one the extract disagrees with the descriptor's field type and the truncation becomes an invalid `llvm.trunc` from `i32` to `i32`.
Read the index type off the descriptor and only adjust the extent width when it actually differs. While here, drop the rank-0 branch that rebuilt the already available element count as an `arith.constant` in the middle of an LLVM lowering.
---
Full diff: https://github.com/llvm/llvm-project/pull/218838.diff
2 Files Affected:
- (modified) mlir/lib/Conversion/MPIToLLVM/MPIToLLVM.cpp (+18-10)
- (modified) mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir (+71-1)
``````````diff
diff --git a/mlir/lib/Conversion/MPIToLLVM/MPIToLLVM.cpp b/mlir/lib/Conversion/MPIToLLVM/MPIToLLVM.cpp
index 70aa735655e18..8179b121ae6cf 100644
--- a/mlir/lib/Conversion/MPIToLLVM/MPIToLLVM.cpp
+++ b/mlir/lib/Conversion/MPIToLLVM/MPIToLLVM.cpp
@@ -15,7 +15,6 @@
#include "mlir/Conversion/MPIToLLVM/MPIToLLVM.h"
#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/DLTI/DLTI.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -56,24 +55,33 @@ std::pair<Value, Value> getRawPtrAndSize(const Location loc,
Value memRef, int64_t rank,
Type elType) {
Type ptrType = LLVM::LLVMPointerType::get(rewriter.getContext());
+ Type i32Type = rewriter.getI32Type();
+ auto descriptorType = cast<LLVM::LLVMStructType>(memRef.getType());
+ // The offset and the sizes of a memref descriptor have the converted index
+ // type, which is not necessarily `i64`. Take it from the descriptor itself.
+ auto indexType = cast<IntegerType>(descriptorType.getBody()[2]);
+
Value dataPtr =
LLVM::ExtractValueOp::create(rewriter, loc, ptrType, memRef, 1);
- Value offset = LLVM::ExtractValueOp::create(rewriter, loc,
- rewriter.getI64Type(), memRef, 2);
+ Value offset =
+ LLVM::ExtractValueOp::create(rewriter, loc, indexType, memRef, 2);
Value resPtr =
LLVM::GEPOp::create(rewriter, loc, ptrType, elType, dataPtr, offset);
- Value size = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
+ Value size = LLVM::ConstantOp::create(rewriter, loc, i32Type,
rewriter.getIndexAttr(1));
- if (cast<LLVM::LLVMStructType>(memRef.getType()).getBody().size() > 3) {
+ if (descriptorType.getBody().size() > 3) {
for (int64_t i = 0; i < rank; ++i) {
Value dim = LLVM::ExtractValueOp::create(rewriter, loc, memRef,
ArrayRef<int64_t>{3, i});
- dim = LLVM::TruncOp::create(rewriter, loc, rewriter.getI32Type(), dim);
- size =
- LLVM::MulOp::create(rewriter, loc, rewriter.getI32Type(), dim, size);
+ // The MPI interface counts elements in an `i32`, so adjust the
+ // index-typed extent to that width. Extents are non-negative, hence the
+ // zero extension.
+ if (indexType.getWidth() > 32)
+ dim = LLVM::TruncOp::create(rewriter, loc, i32Type, dim);
+ else if (indexType.getWidth() < 32)
+ dim = LLVM::ZExtOp::create(rewriter, loc, i32Type, dim);
+ size = LLVM::MulOp::create(rewriter, loc, i32Type, dim, size);
}
- } else {
- size = arith::ConstantIntOp::create(rewriter, loc, 1, 32);
}
return {resPtr, size};
}
diff --git a/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir b/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir
index 73ad2d8f9299f..a8b80e79c00b0 100644
--- a/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir
+++ b/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir
@@ -1,4 +1,6 @@
-// RUN: mlir-opt -split-input-file -convert-to-llvm %s | FileCheck %s
+// `dynamic=true` makes the conversion respect the module's data layout,
+// which is needed to test index types other than the default `i64`.
+// RUN: mlir-opt -split-input-file -convert-to-llvm="dynamic=true" %s | FileCheck %s
// COM: Test MPICH ABI
// CHECK-LABEL: module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">} {
@@ -325,3 +327,71 @@ module attributes {mpi.dlti = #dlti.map<"MPI:Implementation" = "MPICH", "MPI:com
return
}
}
+
+// -----
+
+// COM: Test that an index type that already matches the MPI element count width
+// COM: is used as is, both for the offset and for the extents.
+module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">,
+ dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>} {
+ // CHECK-LABEL: llvm.func @test_send_index32
+ func.func @test_send_index32(%arg0: memref<100xf32>, %rank: i32) {
+ // CHECK: [[v0:%.*]] = llvm.insertvalue {{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i32, array<1 x i32>, array<1 x i32>)>
+ %comm = mpi.comm_world : !mpi.comm
+ // CHECK: [[v1:%.*]] = llvm.extractvalue [[v0]][1] : !llvm.struct<(ptr, ptr, i32, array<1 x i32>, array<1 x i32>)>
+ // CHECK: [[v2:%.*]] = llvm.extractvalue [[v0]][2] : !llvm.struct<(ptr, ptr, i32, array<1 x i32>, array<1 x i32>)>
+ // CHECK: [[v3:%.*]] = llvm.getelementptr [[v1]][[[v2]]] : (!llvm.ptr, i32) -> !llvm.ptr, f32
+ // CHECK: [[v4:%.*]] = llvm.mlir.constant(1 : index) : i32
+ // CHECK: [[v5:%.*]] = llvm.extractvalue [[v0]][3, 0] : !llvm.struct<(ptr, ptr, i32, array<1 x i32>, array<1 x i32>)>
+ // COM: No width adjustment, the extent is already an `i32`.
+ // CHECK-NOT: llvm.trunc
+ // CHECK-NOT: llvm.zext
+ // CHECK: [[v6:%.*]] = llvm.mul [[v5]], [[v4]] : i32
+ // CHECK: llvm.call @MPI_Send([[v3]], [[v6]], {{.*}}) : (!llvm.ptr, i32, i32, i32, i32, i32) -> i32
+ mpi.send(%arg0, %rank, %rank, %comm) : memref<100xf32>, i32, i32
+ return
+ }
+}
+
+// -----
+
+// COM: Test that an index type narrower than the MPI element count is zero
+// COM: extended.
+module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">,
+ dlti.dl_spec = #dlti.dl_spec<index = 16 : i32>} {
+ // CHECK-LABEL: llvm.func @test_send_index16
+ func.func @test_send_index16(%arg0: memref<100xf32>, %rank: i32) {
+ // CHECK: [[v0:%.*]] = llvm.insertvalue {{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i16, array<1 x i16>, array<1 x i16>)>
+ %comm = mpi.comm_world : !mpi.comm
+ // CHECK: [[v1:%.*]] = llvm.extractvalue [[v0]][1] : !llvm.struct<(ptr, ptr, i16, array<1 x i16>, array<1 x i16>)>
+ // CHECK: [[v2:%.*]] = llvm.extractvalue [[v0]][2] : !llvm.struct<(ptr, ptr, i16, array<1 x i16>, array<1 x i16>)>
+ // CHECK: [[v3:%.*]] = llvm.getelementptr [[v1]][[[v2]]] : (!llvm.ptr, i16) -> !llvm.ptr, f32
+ // CHECK: [[v4:%.*]] = llvm.mlir.constant(1 : index) : i32
+ // CHECK: [[v5:%.*]] = llvm.extractvalue [[v0]][3, 0] : !llvm.struct<(ptr, ptr, i16, array<1 x i16>, array<1 x i16>)>
+ // CHECK: [[v6:%.*]] = llvm.zext [[v5]] : i16 to i32
+ // CHECK: [[v7:%.*]] = llvm.mul [[v6]], [[v4]] : i32
+ // CHECK: llvm.call @MPI_Send([[v3]], [[v7]], {{.*}}) : (!llvm.ptr, i32, i32, i32, i32, i32) -> i32
+ mpi.send(%arg0, %rank, %rank, %comm) : memref<100xf32>, i32, i32
+ return
+ }
+}
+
+// -----
+
+// COM: Test that a rank-zero memref, whose descriptor carries no extents, uses
+// COM: the element count of one directly.
+module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">} {
+ // CHECK-LABEL: llvm.func @test_send_rank_zero
+ func.func @test_send_rank_zero(%arg0: memref<f32>, %rank: i32) {
+ // CHECK: [[v0:%.*]] = llvm.insertvalue {{.*}}[2] : !llvm.struct<(ptr, ptr, i64)>
+ %comm = mpi.comm_world : !mpi.comm
+ // CHECK: [[v1:%.*]] = llvm.extractvalue [[v0]][1] : !llvm.struct<(ptr, ptr, i64)>
+ // CHECK: [[v2:%.*]] = llvm.extractvalue [[v0]][2] : !llvm.struct<(ptr, ptr, i64)>
+ // CHECK: [[v3:%.*]] = llvm.getelementptr [[v1]][[[v2]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+ // CHECK: [[v4:%.*]] = llvm.mlir.constant(1 : index) : i32
+ // CHECK-NOT: llvm.mul
+ // CHECK: llvm.call @MPI_Send([[v3]], [[v4]], {{.*}}) : (!llvm.ptr, i32, i32, i32, i32, i32) -> i32
+ mpi.send(%arg0, %rank, %rank, %comm) : memref<f32>, i32, i32
+ return
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/218838
More information about the llvm-branch-commits
mailing list