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 _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
