[mlir] Async: add alwaysinline/noinline attributes to parallel-for functions

[ezhulenev]

Summary

Depends On D119959

+ fix numUnrollableLoops calculation

[mehdi_amini]

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

268

Can you elaborate the comment on why?

"alwaysinline" is a "correctness" kind of attribute, it forces inlining even at O0 in LLVM, is this the case here?

[cota]

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

804–806

IIUC if we had just one loop we'd miss the unroll entirely (since getNumLoops() - 2 == -1, and this patch fixes that.
Is my understanding correct? If so, could you add a small test to show this now working as intended?

[ezhulenev]

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

804–806

Turned out it was WAI, I'll revert this change, and will send a separate chance once I'll figure out what's going with inlining,

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

static ParallelComputeFunction createParallelComputeFunction(
   FunctionType type = computeFuncType.type;
   FuncOp func = FuncOp::create(op.getLoc(),
                                numBlockAlignedInnerLoops > 0
                                    ? "parallel_compute_fn_with_aligned_loops"
                                    : "parallel_compute_fn",
                                type);
   func.setPrivate();
 
+  // Always inline parallel compute function. It gets inlined twice: into the
+  // `async.execute` body and into the main function computing result for the
+  // first block.
+  auto alwaysinline = rewriter.getStringAttr("alwaysinline");
+  func->setAttr("passthrough", rewriter.getArrayAttr({alwaysinline}));
+
   // Insert function into the module symbol table and assign it unique name.
   SymbolTable symbolTable(module);
   symbolTable.insert(func);
   rewriter.getListener()->notifyOperationInserted(func);
 
   // Create function entry block.
   Block *block =
       b.createBlock(&func.getBody(), func.begin(), type.getInputs(),
static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
   inputTypes.push_back(async::GroupType::get(rewriter.getContext()));
   inputTypes.push_back(rewriter.getIndexType()); // add blockStart argument
   inputTypes.append(computeFuncInputTypes.begin(), computeFuncInputTypes.end());
 
   FunctionType type = rewriter.getFunctionType(inputTypes, TypeRange());
   FuncOp func = FuncOp::create(loc, "async_dispatch_fn", type);
   func.setPrivate();
 
+  // TODO: We do not inline async dispatch function, in some cases it leads to
+  // performance regressions because of poor vectorization of compute function.
+  auto noinline = rewriter.getStringAttr("noinline");
+  func->setAttr("passthrough", rewriter.getArrayAttr({noinline}));
+
   // Insert function into the module symbol table and assign it unique name.
   SymbolTable symbolTable(module);
   symbolTable.insert(func);
   rewriter.getListener()->notifyOperationInserted(func);
 
   // Create function entry block.
   Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs(),
                                SmallVector<Location>(type.getNumInputs(), loc));
auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     // than the `maxIterations` value.
     int numUnrollableLoops = 0;
 
     auto getInt = [](IntegerAttr attr) { return attr ? attr.getInt() : 0; };
 
     SmallVector<int64_t> numIterations(op.getNumLoops());
     numIterations.back() = getInt(staticBounds.tripCounts.back());
 
+    // Check if we can potentially unroll the inner-most loop.
+    if (numIterations.back() > 0 && numIterations.back() <= maxIterations)
+      numUnrollableLoops++;
+
+    // Check if we can potentially unroll other loops in the loop nest.
     for (int i = op.getNumLoops() - 2; i >= 0; --i) {
       int64_t tripCount = getInt(staticBounds.tripCounts[i]);
       int64_t innerIterations = numIterations[i + 1];
       numIterations[i] = tripCount * innerIterations;
 
       // Update the number of inner loops that we can potentially unroll.
-      if (innerIterations > 0 && innerIterations <= maxIterations)
+      if (numIterations[i] > 0 && numIterations[i] <= maxIterations)
         numUnrollableLoops++;
     }
 
     Value numWorkerThreadsVal;
     if (numWorkerThreads >= 0)
       numWorkerThreadsVal = b.create<arith::ConstantIndexOp>(numWorkerThreads);
     else
       numWorkerThreadsVal = b.create<async::RuntimeNumWorkerThreadsOp>();

mlir/test/Dialect/Async/async-parallel-for-async-dispatch.mlir

 // CHECK:       scf.for
 // CHECK:         memref.store
 
 // CHECK-LABEL: func private @async_dispatch_fn
 // CHECK-SAME:  (
 // CHECK-SAME:    %[[GROUP:arg0]]: !async.group,
 // CHECK-SAME:    %[[BLOCK_START:arg1]]: index
 // CHECK-SAME:    %[[BLOCK_END:arg2]]: index
-// CHECK-SAME:  )
+// CHECK-SAME:  ) attributes {passthrough = ["noinline"]}
 // CHECK:         %[[C1:.*]] = arith.constant 1 : index
 // CHECK:         %[[C2:.*]] = arith.constant 2 : index
 // CHECK:         scf.while (%[[S0:.*]] = %[[BLOCK_START]],
 // CHECK-SAME:               %[[E0:.*]] = %[[BLOCK_END]])
 // While loop `before` block decides if we need to dispatch more tasks.
 // CHECK:         {
 // CHECK:           %[[DIFF0:.*]] = arith.subi %[[E0]], %[[S0]]
 // CHECK:           %[[COND:.*]] = arith.cmpi sgt, %[[DIFF0]], %[[C1]]

mlir/test/Dialect/Async/async-parallel-for-compute-fn.mlir

 // CHECK-LABEL: func private @parallel_compute_fn(
 // CHECK-SAME:   %[[BLOCK_INDEX:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[BLOCK_SIZE:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[TRIP_COUNT:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?xf32>
-// CHECK-SAME: ) {
+// CHECK-SAME: ) attributes {passthrough = ["alwaysinline"]} {
 // CHECK:        %[[CST:.*]] = arith.constant 1.0{{.*}} : f32
 // CHECK:        scf.for
 // CHECK:          memref.store %[[CST]], %[[MEMREF]]
 
 // -----
 
 // Check that constant loop bound sunk into the parallel compute function.
 

 // CHECK-LABEL: func private @parallel_compute_fn(
 // CHECK-SAME:   %[[BLOCK_INDEX:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[BLOCK_SIZE:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[TRIP_COUNT:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?xf32>
-// CHECK-SAME: ) {
+// CHECK-SAME: ) attributes {passthrough = ["alwaysinline"]} {
 // CHECK:        %[[CSTEP:.*]] = arith.constant 123 : index
 // CHECK-NOT:    %[[STEP]]
 // CHECK:        scf.for %[[I:arg[0-9]+]]
 // CHECK:          %[[TMP:.*]] = arith.muli %[[I]], %[[CSTEP]]
 // CHECK:          %[[IDX:.*]] = arith.addi %[[LB]], %[[TMP]]
 // CHECK:          memref.store
 
 // -----

 // CHECK-SAME:   %[[TRIP_COUNT1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?x10xf32>
-// CHECK-SAME: ) {
+// CHECK-SAME: ) attributes {passthrough = ["alwaysinline"]} {
 // CHECK:        scf.for %[[I:arg[0-9]+]]
 // CHECK:          arith.select
 // CHECK:          scf.for %[[J:arg[0-9]+]]
 // CHECK:          memref.store
 
 // CHECK-LABEL: func private @parallel_compute_fn_with_aligned_loops(
 // CHECK-SAME:   %[[BLOCK_INDEX:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[BLOCK_SIZE:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[TRIP_COUNT0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[TRIP_COUNT1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[LB1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[UB1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP0:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[STEP1:arg[0-9]+]]: index,
 // CHECK-SAME:   %[[MEMREF:arg[0-9]+]]: memref<?x10xf32>
-// CHECK-SAME: ) {
+// CHECK-SAME: ) attributes {passthrough = ["alwaysinline"]} {
 // CHECK:        %[[C0:.*]] = arith.constant 0 : index
 // CHECK:        %[[C1:.*]] = arith.constant 1 : index
 // CHECK:        %[[C10:.*]] = arith.constant 10 : index
 // CHECK:        scf.for %[[I:arg[0-9]+]]
 // CHECK-NOT:      arith.select
 // CHECK:          scf.for %[[J:arg[0-9]+]] = %c0 to %c10 step %c1