[mlir][memref][nfc] push early-exit to earlier

[javedabsar1]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Javed Absar (javedabsar1)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/140730.diff

1 Files Affected:

• (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+4-4)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 7f45904fab7e1..d78de25f3e0ff 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -1711,6 +1711,10 @@ struct ViewOpLowering : public ConvertOpToLLVMPattern<memref::ViewOp> {
     MemRefDescriptor sourceMemRef(adaptor.getSource());
     auto targetMemRef = MemRefDescriptor::poison(rewriter, loc, targetDescTy);
 
+    // Early exit for 0-D corner case.
+    if (viewMemRefType.getRank() == 0)
+      return rewriter.replaceOp(viewOp, {targetMemRef}), success();
+
     // Field 1: Copy the allocated pointer, used for malloc/free.
     Value allocatedPtr = sourceMemRef.allocatedPtr(rewriter, loc);
     auto srcMemRefType = cast<MemRefType>(viewOp.getSource().getType());
@@ -1733,10 +1737,6 @@ struct ViewOpLowering : public ConvertOpToLLVMPattern<memref::ViewOp> {
         rewriter, loc,
         createIndexAttrConstant(rewriter, loc, indexType, offset));
 
-    // Early exit for 0-D corner case.
-    if (viewMemRefType.getRank() == 0)
-      return rewriter.replaceOp(viewOp, {targetMemRef}), success();
-
     // Fields 4 and 5: Update sizes and strides.
     Value stride = nullptr, nextSize = nullptr;
     for (int i = viewMemRefType.getRank() - 1; i >= 0; --i) {

[banach-space]

I am not familiar with this code, but the change is quite straightforward and a clear improvement. Thanks, LGTM!

[Garra1980]

What was the rationale behind this change? I see that now for the following snippet

func.func @test() attributes {llvm.emit_c_interface} {
    %c0 = arith.constant 0 : index

    %memref_bf16 = gpu.alloc  host_shared () : memref<bf16>

    %memref_i8 = gpu.alloc  host_shared () : memref<2xi8>
    %memref_view_bf16 = memref.view %memref_i8[%c0][] : memref<2xi8> to memref<bf16>
    memref.copy %memref_bf16, %memref_view_bf16 : memref<bf16> to memref<bf16>


    return
  }

after --finalize-memref-to-llvm IR looks as

    %memref_0 = gpu.alloc  host_shared () : memref<2xi8>
    %1 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64)>
    %2 = llvm.mlir.constant(1 : index) : i64
    %3 = llvm.mlir.zero : !llvm.ptr
    %4 = llvm.getelementptr %3[1] : (!llvm.ptr) -> !llvm.ptr, bf16
    %5 = llvm.ptrtoint %4 : !llvm.ptr to i64
    %6 = llvm.mul %2, %5 : i64
    %7 = llvm.extractvalue %0[1] : !llvm.struct<(ptr, ptr, i64)>
    %8 = llvm.extractvalue %0[2] : !llvm.struct<(ptr, ptr, i64)>
    %9 = llvm.getelementptr %7[%8] : (!llvm.ptr, i64) -> !llvm.ptr, bf16
    %10 = llvm.extractvalue %1[1] : !llvm.struct<(ptr, ptr, i64)>
    %11 = llvm.extractvalue %1[2] : !llvm.struct<(ptr, ptr, i64)>
    %12 = llvm.getelementptr %10[%11] : (!llvm.ptr, i64) -> !llvm.ptr, bf16
    "llvm.intr.memcpy"(%12, %9, %6) <{isVolatile = false}> : (!llvm.ptr, !llvm.ptr, i64) -> ()

so no flow from %12 to memref_i8 which seems to be wrong

[banach-space]

Hey @Garra1980 , thanks for posting this and apologies if you're experiencing issues due to this PR.

What was the rationale behind this change?

Early exits are encouraged in LLVM and this looked like (and still does) like an NFC.

no flow from %12 to memref_i8 which seems to be wrong

It's not quite clear to me what is wrong 😅 Could you compare before and after IR? If after is indeed incorrect, we should include your example as a test and revert this change.

[Garra1980]

It's not quite clear to me what is wrong 😅 Could you compare before and after IR? If after is indeed incorrect, we should include your example as a test and revert this change.

Ah, my bad! You're right, before it was like this:

    %memref_0 = gpu.alloc  host_shared () : memref<2xi8>
    %2 = builtin.unrealized_conversion_cast %memref_0 : memref<2xi8> to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
    %3 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64)>
    %4 = llvm.extractvalue %2[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
    %5 = llvm.insertvalue %4, %3[0] : !llvm.struct<(ptr, ptr, i64)> 
    %6 = llvm.extractvalue %2[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
    %7 = llvm.getelementptr %6[%0] : (!llvm.ptr, i64) -> !llvm.ptr, i8
    %8 = llvm.insertvalue %7, %5[1] : !llvm.struct<(ptr, ptr, i64)> 
    %9 = llvm.mlir.constant(0 : index) : i64
    %10 = llvm.insertvalue %9, %8[2] : !llvm.struct<(ptr, ptr, i64)> 
    %11 = llvm.mlir.constant(1 : index) : i64
    %12 = llvm.mlir.zero : !llvm.ptr
    %13 = llvm.getelementptr %12[1] : (!llvm.ptr) -> !llvm.ptr, bf16
    %14 = llvm.ptrtoint %13 : !llvm.ptr to i64
    %15 = llvm.mul %11, %14 : i64
    %16 = llvm.extractvalue %1[1] : !llvm.struct<(ptr, ptr, i64)> 
    %17 = llvm.extractvalue %1[2] : !llvm.struct<(ptr, ptr, i64)> 
    %18 = llvm.getelementptr %16[%17] : (!llvm.ptr, i64) -> !llvm.ptr, bf16
    %19 = llvm.extractvalue %10[1] : !llvm.struct<(ptr, ptr, i64)> 
    %20 = llvm.extractvalue %10[2] : !llvm.struct<(ptr, ptr, i64)> 
    %21 = llvm.getelementptr %19[%20] : (!llvm.ptr, i64) -> !llvm.ptr, bf16
    "llvm.intr.memcpy"(%21, %18, %15) <{isVolatile = false}> : (!llvm.ptr, !llvm.ptr, i64) -> ()

[banach-space]

Thanks @Garra1980 , here's a revert with a minimal test:

• [mlir][memref] Revert #140730 #141406