[MLIR][AFFINE] Handling possible side-effects during fusion #13

SwapnilGhanshyala · 2024-01-23T08:16:18Z

Check if the forOps have any possible side-effects and handle them during MDG formation. This commit checks if forOp contains a memoryEffect operation, which is a possible side-effects and thus, fusion is not possible. In this commit we modifify the dependency graph initialization to add dependency edges between nodes with side-effecting ops and any other node that may be affected by it (basically node that is accessing any memref).

Also, modified test loop-fusion-4.mlir to add a test case.

Issue: maximal loop fusion was not handling side-effecting ops like function calls in the loops.
Original Pipeline: Loops with possible side effects are classified as sequential loops. If the sibling was sequential, it led to a nesting of the loops instead of fusion. In standard fusion (non-maximal), the strategy would have been rejected based on performance heuristics, and another iteration of bounds calculation would start. However, in maximal-fusion, this strategy would not be rejected and hence loops with possible side-effects would be fused.
Solution: MDG initialization needed modification to handle nodes with memoryEffect ops. Edges are now added between nodes with side-effects and nodes that might get impacted.

INPUT:

  func.func @testing2(%arg0: memref<10xf32>) {
    affine.for %i = 0 to 10 {
      %0 = affine.load %arg0[%i] : memref<10xf32>
      llvm.call @g():()->()
    }
    affine.for %j = 0 to 10 {
      %0 = affine.load %arg0[%j] : memref<10xf32>
      llvm.call @f():()->()
    }
    return
  }
  llvm.func external @f()->()
  llvm.func external @g()->()
}

COMMAND: mlir-opt -pass-pipeline='builtin.module(func.func(affine-loop-fusion{fusion-maximal}))'

PREVIOUS OUTPUT:

  func.func @testing2(%arg0: memref<10xf32>) {
    affine.for %arg1 = 0 to 10 {
      affine.for %arg2 = 0 to 10 {
        %1 = affine.load %arg0[%arg2] : memref<10xf32>
        llvm.call @g() : () -> ()
      }
      %0 = affine.load %arg0[%arg1] : memref<10xf32>
      llvm.call @f() : () -> ()
    }
    return
  }
  llvm.func @f()
  llvm.func @g()
}

NEW OUTPUT:

  func.func @testing2(%arg0: memref<10xf32>) {
    affine.for %arg1 = 0 to 10 {
      %0 = affine.load %arg0[%arg1] : memref<10xf32>
      llvm.call @g() : () -> ()
    }
    affine.for %arg1 = 0 to 10 {
      %0 = affine.load %arg0[%arg1] : memref<10xf32>
      llvm.call @f() : () -> ()
    }
    return
  }
  llvm.func @f()
  llvm.func @g()
}

SwapnilGhanshyala · 2024-01-23T08:17:16Z

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

Modified the LoopNestStateCollector structure to hold memoryEffectOps.

SwapnilGhanshyala · 2024-01-23T08:17:30Z

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

Modified the Node struct of MDG to List memoryEffectOps.

SwapnilGhanshyala · 2024-01-23T08:18:07Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

If the op is not memoryEffectFree, LoopStateCollector should add this op to the list of MemoryEffectOps

Is the above filter enough? Should it be else if (!mlir::isMemoryEffectFree(op) && !isa<mlir::memref::AllocOp, mlir::memref::DeallocOp>(op)) , instead of the above.

if only check !mlir::ismemoryEffectFree(op), it conservatively includes alloc and dealloc ops as well. Because of this inclusion code like below fail to fuse successfully.

module { func.func @testing2(%arg0: memref<10xf32>, %arg2 : memref <10xf32>) { %cst = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 2.000000e+00 : f32 %1 = memref.alloc() : memref<12x12xf32> affine.for %arg1 = 0 to 10 { %0 = affine.load %arg0[%arg1] : memref<10xf32> } affine.for %arg1 = 0 to 10 { %0 = affine.load %arg0[%arg1] : memref<10xf32> memref.dealloc %1 : memref<12x12xf32> } return } }

memref.dealloc and memref.alloc can be better handled in the MDG construction by tracking the memref they access similarly to load and store ops.

@bondhugula , current commit conservatively treats alloc and deallocate as having unknown side-effects (as stated in the above comments)
Will fix the above issue by modifying MDG init to handle alloc/dealloc better and update pull request with it.

SwapnilGhanshyala · 2024-01-23T08:18:21Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

This vector holds the for-nodes with potential memory effects in the order of traversal

affine.for nodes?

SwapnilGhanshyala · 2024-01-23T08:18:56Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

When a memory reference is accessed for the first time, all the for-nodes with possible memory effects that precede it are added to its access list.

SwapnilGhanshyala · 2024-01-23T08:19:21Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

If the current for-node exhibits memory effects, it is appended to the sequential list of 'for-nodes' that have memory effects.

SwapnilGhanshyala · 2024-01-23T08:19:39Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

An edge is introduced between the srcNode and dstNode if either has a memory effect or a store op.
The current algorithm tracks memoryEffectOps which are within a loop body. Top-level ops with memoryEffect are handled separately for fusion validity checks.

Please add code comments as suitable instead of github comments.

SwapnilGhanshyala · 2024-01-23T08:19:56Z

mlir/test/Dialect/Affine/loop-fusion-4.mlir

A small test demonstrating the handling of side-effects when fusing loops with fusion-maximal

SwapnilGhanshyala · 2024-01-23T08:21:30Z

Issues addressed here:

Added code comments

SwapnilGhanshyala · 2024-01-27T08:45:43Z

Hi @bondhugula , a gentle reminder for the code review.

bondhugula · 2024-01-27T09:31:37Z

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

Missing code comment here on what memory effect ops are. Load and store also have memory effects. Add doc comment for this variable.

bondhugula · 2024-01-27T09:31:54Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

All comments have to end with a full stop. LLVM style.

bondhugula · 2024-01-27T09:32:13Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

llvm::is_contained

bondhugula · 2024-01-27T09:32:51Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

No need of parens.

bondhugula · 2024-01-27T09:34:31Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

You are combining ops with all memory effects but you need those with write effects below. See below.

Can you respond to this comment?

Currently, we are considering side-effects as unknown entities that could be potentially read or write ops. However, if we consider side-effects with only write effect. we might lose out on potential WAR dependency in the MDGs.

See comment response #13 (comment)

bondhugula · 2024-01-27T09:34:42Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

Variable name will be inaccurate now. dstMayHaveStore.

Moreover, you could have only kept write memory effects here.

SwapnilGhanshyala · 2024-01-30T07:39:05Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

Filtering out alloc and dealloc ops as well. Handling of memref:: alloc/dealloc ops needs to be addressed in a separate as a separate issue to make this commit less complicated.

Here is an example of an existing issue in handling alloc. The following loop does not merge if %alloc is memref.alloc.

func.func @should_not_fuse_loops_with_side_effects() { %alloc = memref.alloc() : memref<10xf32> affine.for %arg0 = 0 to 10 { %0 = affine.load %alloc[%arg0] : memref<10xf32> } affine.for %arg0 = 0 to 10 { %0 = affine.load %alloc[%arg0] : memref<10xf32> } return }

SwapnilGhanshyala · 2024-01-30T07:55:26Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

The check on the destination cannot be just for writes. Otherwise, it will miss out on handling cases where srcNode has load op and dstNode has both load and a memoryEffectOp(which has unknown side-effects).

The following sibling loops below will merge, whereas they shouldn't, if we remove the !dstNode->memEffects.empty() check:

func.func @should_not_fuse_loops_with_side_effects(%arg0:memref<10xf32>) { affine.for %i = 0 to 10 { %0 = affine.load %arg0[%i] : memref<10xf32> } affine.for %j = 0 to 10 { %0 = affine.load %arg0[%j] : memref<10xf32> "foo"():()->() } return }

Also changed variable names to a better representation dstMayHaveMemEffOrStore

SwapnilGhanshyala

@bondhugula , I have addressed the previous comments and added some new ones.

SwapnilGhanshyala · 2024-01-31T08:19:10Z

Hi @bondhugula, a gentle reminder for review.

bondhugula · 2024-01-31T08:51:22Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

Please don't hardcode alloc, dealloc; instead check for Alloc and Free effects.

Corrected in the latest commit.

Check if the forOps have any possible side-effects and handle them during MDG formation. This commit checks if forOp contains a memoryEffect operation, which is a possible side-effects and thus, fusion is not possible. In this commit we modifify the dependency graph initialization to add dependency edges between nodes with side-effecting ops and any other node that may be affected by it (basically node that is accessing any memref). Also, modified test loop-fusion-4.mlir to add a test case.

SwapnilGhanshyala · 2024-02-02T08:05:17Z

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

    else if (isa<AffineReadOpInterface>(op))
      loadOpInsts.push_back(op);
    else if (isa<AffineWriteOpInterface>(op))
      storeOpInsts.push_back(op);


Filtering out ops with alloc and free effects, reason in #13 (comment)

SwapnilGhanshyala · 2024-02-12T08:49:46Z

Hi @bondhugula , gentle reminder for review.

SwapnilGhanshyala self-assigned this Jan 23, 2024

SwapnilGhanshyala commented Jan 23, 2024

View reviewed changes

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h Outdated

Copy link

Author

SwapnilGhanshyala Jan 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Modified the Node struct of MDG to List memoryEffectOps.

SwapnilGhanshyala commented Jan 23, 2024

View reviewed changes

SwapnilGhanshyala requested a review from bondhugula January 23, 2024 08:21

bondhugula reviewed Jan 27, 2024

View reviewed changes

SwapnilGhanshyala force-pushed the side-effect-2 branch from 2d756d0 to 2f33c01 Compare January 30, 2024 07:34

SwapnilGhanshyala commented Jan 30, 2024

View reviewed changes

SwapnilGhanshyala requested a review from bondhugula January 30, 2024 07:58

bondhugula reviewed Jan 31, 2024

View reviewed changes

SwapnilGhanshyala force-pushed the side-effect-2 branch from 2f33c01 to c74638d Compare February 2, 2024 08:00

SwapnilGhanshyala commented Feb 2, 2024

View reviewed changes

SwapnilGhanshyala requested a review from bondhugula February 2, 2024 08:11

[MLIR][AFFINE] Handling possible side-effects during fusion #13

Are you sure you want to change the base?

[MLIR][AFFINE] Handling possible side-effects during fusion #13

Uh oh!

Conversation

SwapnilGhanshyala commented Jan 23, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala Jan 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala Jan 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala commented Jan 23, 2024

Uh oh!

SwapnilGhanshyala commented Jan 27, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala Jan 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala Jan 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala Jan 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala Jan 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SwapnilGhanshyala left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

SwapnilGhanshyala Jan 29, 2024 •

edited

Loading

SwapnilGhanshyala Jan 29, 2024 •

edited

Loading

SwapnilGhanshyala Jan 30, 2024 •

edited

Loading

SwapnilGhanshyala Jan 30, 2024 •

edited

Loading

SwapnilGhanshyala Jan 30, 2024 •

edited

Loading

SwapnilGhanshyala Jan 30, 2024 •

edited

Loading

SwapnilGhanshyala left a comment •

edited

Loading