[P1] Fuse constant into dependent ops

[tancheng]

Only one constant can be folded into a given operation.

[ShangkunLi]

We can fuse each constant and data_mov into dependent ops and add corresponding attributes used for register configuration.

[ShangkunLi]

Provide a prototype in #102.

[tancheng]

Hi @ShangkunLi, are you on this? I can't remember what type of fusion you would go first. I think for:

dataflow/test/neura/ctrl/branch_for.mlir

Lines 154 to 186 in 7f845d6

	// FUSE: func.func @loop_test() -> f32 attributes {accelerator = "neura"} {
	// FUSE-NEXT: %0 = "neura.grant_once"() <{constant_value = 10 : i64}> : () -> !neura.data<i64, i1>
	// FUSE-NEXT: %1 = "neura.grant_once"() <{constant_value = 0 : i64}> : () -> !neura.data<i64, i1>
	// FUSE-NEXT: %2 = "neura.grant_once"() <{constant_value = 1 : i64}> : () -> !neura.data<i64, i1>
	// FUSE-NEXT: %3 = "neura.grant_once"() <{constant_value = 3.000000e+00 : f32}> : () -> !neura.data<f32, i1>
	// FUSE-NEXT: %4 = "neura.grant_once"() <{constant_value = 0.000000e+00 : f32}> : () -> !neura.data<f32, i1>
	// FUSE-NEXT: %5 = neura.reserve : !neura.data<i64, i1>
	// FUSE-NEXT: %6 = "neura.phi"(%5, %0) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
	// FUSE-NEXT: %7 = neura.reserve : !neura.data<i64, i1>
	// FUSE-NEXT: %8 = "neura.phi"(%7, %2) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
	// FUSE-NEXT: %9 = neura.reserve : !neura.data<f32, i1>
	// FUSE-NEXT: %10 = "neura.phi"(%9, %3) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
	// FUSE-NEXT: %11 = neura.reserve : !neura.data<f32, i1>
	// FUSE-NEXT: %12 = "neura.phi"(%11, %4) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
	// FUSE-NEXT: %13 = neura.reserve : !neura.data<i64, i1>
	// FUSE-NEXT: %14 = "neura.phi"(%13, %1) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
	// FUSE-NEXT: %15 = "neura.fadd"(%12, %10) : (!neura.data<f32, i1>, !neura.data<f32, i1>) -> !neura.data<f32, i1>
	// FUSE-NEXT: %16 = "neura.add"(%14, %8) : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i64, i1>
	// FUSE-NEXT: %17 = "neura.icmp"(%16, %6) <{cmpType = "slt"}> : (!neura.data<i64, i1>, !neura.data<i64, i1>) -> !neura.data<i1, i1>
	// FUSE-NEXT: %18 = neura.grant_predicate %16, %17 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
	// FUSE-NEXT: neura.ctrl_mov %18 -> %13 : !neura.data<i64, i1> !neura.data<i64, i1>
	// FUSE-NEXT: %19 = neura.grant_predicate %15, %17 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
	// FUSE-NEXT: neura.ctrl_mov %19 -> %11 : !neura.data<f32, i1> !neura.data<f32, i1>
	// FUSE-NEXT: %20 = neura.grant_predicate %10, %17 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
	// FUSE-NEXT: neura.ctrl_mov %20 -> %9 : !neura.data<f32, i1> !neura.data<f32, i1>
	// FUSE-NEXT: %21 = neura.grant_predicate %8, %17 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
	// FUSE-NEXT: neura.ctrl_mov %21 -> %7 : !neura.data<i64, i1> !neura.data<i64, i1>
	// FUSE-NEXT: %22 = neura.grant_predicate %6, %17 : !neura.data<i64, i1>, !neura.data<i1, i1> -> !neura.data<i64, i1>
	// FUSE-NEXT: neura.ctrl_mov %22 -> %5 : !neura.data<i64, i1> !neura.data<i64, i1>
	// FUSE-NEXT: %23 = "neura.not"(%17) : (!neura.data<i1, i1>) -> !neura.data<i1, i1>
	// FUSE-NEXT: %24 = neura.grant_predicate %15, %23 : !neura.data<f32, i1>, !neura.data<i1, i1> -> !neura.data<f32, i1>
	// FUSE-NEXT: "neura.return"(%24) : (!neura.data<f32, i1>) -> ()
	// FUSE-NEXT: }

We can fuse grant_once()<const> further into phi and cmp, to make them be phi_once()<const> and cmp<const>, respectively, the number of critical recurrent cycles would significantly decrease. It should be easier to be implemented based on your current implementation here:

dataflow/lib/NeuraDialect/Transforms/FoldConstantPass.cpp

Lines 23 to 44 in 7f845d6

	for (auto user : constant_op->getUsers()) {
	llvm::errs() << "Checking use: " << *user << "\n";
	if (isa<neura::GrantOnceOp>(user) || isa<neura::GrantAlwaysOp>(user)) {
	if (neura::GrantOnceOp grant_once_op =
	dyn_cast<neura::GrantOnceOp>(user)) {
	auto new_grant_once_op = rewriter.create<neura::GrantOnceOp>(
	grant_once_op.getLoc(), grant_once_op.getResult().getType(),
	/*value=*/nullptr, constant_op->getAttr("value"));
	// Replaces the original constant operation with the new one.
	rewriter.replaceOp(grant_once_op, new_grant_once_op);
	made_change = true;
	} else if (neura::GrantAlwaysOp grant_always_op =
	dyn_cast<neura::GrantAlwaysOp>(user)) {
	auto new_grant_always_op = rewriter.create<neura::GrantAlwaysOp>(
	grant_always_op.getLoc(), grant_always_op.getResult().getType(),
	/*value=*/nullptr, constant_op->getAttr("value"));
	// Replaces the original constant operation with the new one.
	rewriter.replaceOp(grant_always_op, new_grant_always_op);
	made_change = true;
	}
	}
	}

Two things we need to discuss though:

• Is it an appropriate design that we give const as an attribute rather than an operand for an IR?
• Is it robust to always fold all predicate_once<const> into the corresponding dependent op?
• It is not clear to me how to easily fold const with the cmp, it is straightforward to me when hand-write the kernel, but not that straightforward after we canonicalizing the BB arguments
  • I still think BB canonicalization is the right way to go for robustness of complex ctrl-flow and nested for loops, but maybe we can fold const before BB canonicalization?

[ShangkunLi]

I am working on the paper writing today. I will begin to handle operation fusion tomorrow.

We can fuse grant_once()<const> further into phi and cmp, to make them be phi_once()<const> and cmp<const>, respectively, the number of critical recurrent cycles would significantly decrease. It should be easier to be implemented based on your current implementation here:

Fusing grant_once into phi would not decrease the number of recurrence cycles. Because there will always be a reserve and a ctrl_mov for a phi. My plan is to detect the two sources of a phi, if they are from the same source, then we can remove this phi and keep grant_predicate(grant_once, %cond). This can remove the recurrence cycle.

Fusing grant_once into cmp is straightforward, and I will do it. I think if the source value of the bound of cmp is a constant, we can directly fuse it into the cmp.

Is it an appropriate design that we give const as an attribute rather than an operand for an IR?

If we choose to keep constant ops, what else operations can we fuse?

Is it robust to always fold all predicate_once into the corresponding dependent op?

The --fold-constant pass will keep the constant ops if they are directly used in other ops.

• I still think BB canonicalization is the right way to go for robustness of complex ctrl-flow and nested for loops, but maybe we can fold const before BB canonicalization?

I will try to fold the constant first and then do the following transforms.

[tancheng]

Thanks @ShangkunLi, if we fold first then transform, maybe we don't need to care about "Because there will always be a reserve and a ctrl_mov for a phi. My plan is to detect the two sources of a phi, if they are from the same source, then we can remove this phi and keep grant_predicate(grant_once, %cond). This can remove the recurrence cycle."

Such folding before transform can be the first PR. And I think the implementation can be added into the current fold-const-pass? And make it be called twice (before and after ctrl-to-data-pass), though the second time calling might be not necessary any more.

[ShangkunLi]

Thanks @ShangkunLi, if we fold first then transform, maybe we don't need to care about "Because there will always be a reserve and a ctrl_mov for a phi. My plan is to detect the two sources of a phi, if they are from the same source, then we can remove this phi and keep grant_predicate(grant_once, %cond). This can remove the recurrence cycle."

Such folding before transform can be the first PR. And I think the implementation can be added into the current fold-const-pass? And make it be called twice (before and after ctrl-to-data-pass), though the second time calling might be not necessary any more.

Sure~ I will try it. But if we choose to fold constant before ctrl-to-data-pass, we still need to handle phi elimination because phi and ctrl_mov are inserted during ctrl-to-data-pass.

[ShangkunLi]

Hi~ @tancheng,

I think a better way to reduce the number of ops would be to fuse all the operations with their constant operands before canonicalize-live-in, if any. Like, fusing %3 with %10, %2 with %11, %1 with %8, etc.

And the predication can be propagated properly after canonicalize-live-in and transform-ctrl-to-data-flow. WDYT?

module {
  func.func @_Z11simple_loopPiS_(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {accelerator = "neura", llvm.linkage = #llvm.linkage<external>} {
    %0 = "neura.constant"() <{predicate = true, value = 1 : index}> : () -> index
    %1 = "neura.constant"() <{predicate = true, value = 128 : index}> : () -> index
    %2 = "neura.constant"() <{predicate = true, value = 1 : i32}> : () -> i32
    %3 = "neura.constant"() <{predicate = true, value = 2 : i32}> : () -> i32
    %4 = "neura.constant"() <{predicate = true, value = 0 : index}> : () -> index
    %5 = "neura.cast"(%4) <{cast_type = "index_to_int"}> : (index) -> i64
    neura.br %5 : i64 to ^bb1
  ^bb1(%6: i64):  // 2 preds: ^bb0, ^bb2
    %7 = "neura.cast"(%6) <{cast_type = "int_to_index"}> : (i64) -> index
    %8 = "neura.icmp"(%7, %1) <{cmpType = "slt"}> : (index, index) -> i1
    neura.cond_br %8 : i1 then to ^bb2 else to ^bb3
  ^bb2:  // pred: ^bb1
    %9 = neura.load_indexed %arg0[%7 : index] memref<?xi32> : i32
    %10 = "neura.mul"(%9, %3) : (i32, i32) -> i32
    %11 = "neura.add"(%10, %2) : (i32, i32) -> i32
    neura.store_indexed %11 to %arg1[%7 : index] memref<?xi32> : i32
    %12 = "neura.add"(%7, %0) : (index, index) -> index
    %13 = "neura.cast"(%12) <{cast_type = "index_to_int"}> : (index) -> i64
    neura.br %13 : i64 to ^bb1
  ^bb3:  // pred: ^bb1
    "neura.return"() : () -> ()
  }
}

[tancheng]

That's exactly why I mention "maybe we don't need to care about xxxx". Plz go ahead.

• I am not sure whether we should simply allow all const fusion with complicated BB domination dependency, all fusions in your provided test make sense. I am not able to come up a counter-example right now...
• We may first just fuse rhs const as const attribute, or you think either lhs const and rhs const can be done, anyways we don't have a case both rhs and lhs are const