[async] Get the number of worker threads from the runtime.

[ftynse]

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

543

Why is this necessary? The return type is always index, ODS should not how to build it. If something is not working without this, it is probably a missing ODS functionality, we shouldn't need to write such boilerplate.

544–547

Nit: put double indentation for function arguments, otherwise it's hard to differentiate them from the body.

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

799

Nit: drop trivial braces.

820

Please avoid magic numbers.

823

const &

831

TypeRange should be implicitly constructible from a single type.

840–841

.getResult(0) looks shorter.

844–855

Please avoid nested calls to Builder::create. Beyond being harder to read, it leads to a platform-dependent order of operations since the order of evaluation for function call arguments is implementation-defined in C++. I suspect this is the reason why the test is broken on Windows (makeSwitch creates operations and so does the other argument of MulFOp).

mlir/test/Dialect/Async/async-parallel-for-num-worker-threads.mlir

12

Nit: why do we even emit this conditional if we know the condition?

15

Please avoid matching SSA value names, even for constants.

[ezhulenev]

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

541

I'd prefer to see : type($result) here, even if ODS can infer it, explicit typing makes it easier to read, and if ODS can't do it now then inferReturnTypes is too much.

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

128

nit: empty line before the next comment

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

799

For numWorkerThreads == 0 this pass should be no-op, otherwise I assume it will crash with some division by zero

811

Can you send this change in a separate PR, not to mix it with straightforward async runtime change

811

Ah, just realized that numWorkerThreads is no longer a constant.

[ezhulenev]

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

810

Can you include simplified version of the old code in comments, because parsing what's going on here from builders is too hard.

819

Here and in all other labmdas capture can be just [&] because they do not leave the scope.

[ezhulenev]

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

823

Instead of generating nested scf.if you can do that with a short number of cmpi and select, I think it will be easier to match in tests and read in mlir dumps, and no need to to recursive builders, just a loop over brackets will do the work

%threads = async.runtime.num_threads

%init = constant 8.0 : float // (0, 4]

%p4 = cmpi sgt %threads %c4     // num_threads > 4
%p8 = cmpi sgt %threads %c8     // num_threads > 8
%p16 = cmpi sgt %threads %c16 // num_threads > 16
%p32 = cmpi sgt %threads %c32 // num_threads > 32
%p64 = cmpi sgt %threads %c64 // num_threads > 64

%s0 = select %p4   %c4.0  %c8.0
%s1 = select %p8   %c2.0  %s0
%s2 = select %p16 %c1.0  %s1
%s3 = select %p32 %c0.8  %s2
%scale = select %p64 %c0.6  %s3

[bakhtiyarneyman]

Address comments.

[bakhtiyarneyman]

Revert changes to Bazel file because it doesn't belong in this CL.

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

541

Done.

543

I was confused about interfaces: I didn't realize that merely importing the interface module is sufficient for it to take effect. I thought that is is necessary to use it as a trait on an op.

Removed.

544–547

Obsolete.

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

799

Definition of maxComputeBlock prevents that. I'm not changing logic in this regard. I suppose an early exit is possible, but I'd argue it belongs in a separate PR.

810

Done.

811

I'm assuming this is withdrawn.

819

Obsolete. FWIW, makeSwitch needed to be passed by copy, because it was supposed to replace past self, so I opted for maximum explicitness. But it's moot point now.

823

Great!

831

Obsolete.

840–841

Obsolete.

844–855

Done.

[ezhulenev]

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

824

You can include "mlir/Support/LLVM.h" to bring APFloat into mlir namespace (for consistency with other mlir files), also a separate makeFloat doesn't make much sense now, its used just twice, and "creating custom mini DSLs" is frowned upon in MLIR

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

441

auto *runtime = getDefaultAsyncRuntime();
  runtime->getThreadPool()-> <get the real number of worker threads> ``` (snippet from `mlirAsyncRuntimeExecute`

[ezhulenev]

Can you please also add tests cases to tests (benchmarks?) in this folder: https://github.com/llvm/llvm-project/tree/main/mlir/test/Integration/Dialect/Async/CPU

[bakhtiyarneyman]

Address comments.

[bakhtiyarneyman]

PTAL.

[ezhulenev]

I get this error locally:

AILED: lib/libMLIRAsync.so.14git
: && /usr/bin/c++ -fPIC -fPIC -fno-semantic-interposition -fvisibility-inlines-hidden -Werror=date-time -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wno-missing-field-initializers -pedantic -Wno-long-long -Wimplicit-fallthrough -Wno-maybe-uninitialized -Wno-class-memaccess -Wno-redundant-move -Wno-pessimizing-move -Wno-noexcept-type -Wdelete-non-virtual-dtor -Wsuggest-override -Wno-comment -Wmisleading-indentation -fdiagnostics-color -g  -Wl,-z,defs -Wl,-z,nodelete   -Wl,-rpath-link,/usr/local/google/home/ezhulenev/projects/llvm-project/build/./lib -shared -Wl,-soname,libMLIRAsync.so.14git -o lib/libMLIRAsync.so.14git tools/mlir/lib/Dialect/Async/IR/CMakeFiles/obj.MLIRAsync.dir/Async.cpp.o  -Wl,-rpath,"\$ORIGIN/../lib"  lib/libMLIRDialect.so.14git  lib/libMLIRIR.so.14git  lib/libMLIRSupport.so.14git  -lpthread  lib/libLLVMSupport.so.14git  -Wl,-rpath-link,/usr/local/google/home/ezhulenev/projects/llvm-project/build/lib && :
/usr/bin/ld: tools/mlir/lib/Dialect/Async/IR/CMakeFiles/obj.MLIRAsync.dir/Async.cpp.o: in function `mlir::detail::InferTypeOpInterfaceTrait<mlir::async::AddToGroupOp>::verifyTrait(mlir::Operation*)':
/usr/local/google/home/ezhulenev/projects/llvm-project/build/tools/mlir/include/mlir/Interfaces/InferTypeOpInterface.h.inc:141: undefined reference to `mlir::detail::verifyInferredResultTypes(mlir::Operation*)'

I think CMakeLists has to be updated to be consistent with BUILD.bazel, though it seems that InferTypeOpInterface is no longer used and can be deleted

[bakhtiyarneyman]

Update relevant CMakeList.txt.

mlir/include/mlir/Dialect/Async/IR/Async.h

 #include "mlir/Dialect/Async/IR/AsyncTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Async Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Async/IR/AsyncOpsDialect.h.inc"
 

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

 //===----------------------------------------------------------------------===//
 
 #ifndef ASYNC_OPS
 #define ASYNC_OPS
 
 include "mlir/Dialect/Async/IR/AsyncDialect.td"
 include "mlir/Dialect/Async/IR/AsyncTypes.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // Async op definitions
 //===----------------------------------------------------------------------===//
 
 // Base class for the operation in this dialect
 class Async_Op<string mnemonic, list<OpTrait> traits = []> :
def Async_RuntimeDropRefOp : Async_Op<"runtime.drop_ref"> {
   let arguments = (ins Async_AnyAsyncType:$operand,
                        Confined<I64Attr, [IntPositive]>:$count);
 
   let assemblyFormat = [{
     $operand attr-dict `:` type($operand)
   }];
 }
 
+def Async_RuntimeNumWorkerThreadsOp : Async_Op<"runtime.num_worker_threads", [InferTypeOpInterface]> {
+  let summary = "gets the number of threads in the threadpool from the runtime";
+  let description = [{
+    The `async.runtime.num_worker_threads` operation gets the number of threads
+    in the threadpool from the runtime.
+  }];
+
+  let results = (outs Index:$result);
+  let assemblyFormat = "attr-dict";
+  let extraClassDeclaration = [{
+    static ::mlir::LogicalResult inferReturnTypes(
+      ::mlir::MLIRContext *context, ::llvm::Optional<::mlir::Location> location,
+      ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
+      ::mlir::RegionRange regions,
+      ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+      inferredReturnTypes.resize(1);
+      ::mlir::Builder odsBuilder(context);
+      inferredReturnTypes[0] = odsBuilder.getIndexType();
+      return ::mlir::success();
+    }
+  }];
+}
+
 #endif // ASYNC_OPS

mlir/include/mlir/Dialect/Async/Passes.td

let options = [
     Option<"asyncDispatch", "async-dispatch",
       "bool", /*default=*/"true",
       "Dispatch async compute tasks using recursive work splitting. If `false` "
       "async compute tasks will be launched using simple for loop in the "
       "caller thread.">,
 
     Option<"numWorkerThreads", "num-workers",
       "int32_t", /*default=*/"8",
-      "The number of available workers to execute async operations.">,
+      "The number of available workers to execute async operations. If `-1` "
+      "the value will be retrieved from the runtime.">,
 
     Option<"minTaskSize", "min-task-size",
       "int32_t", /*default=*/"1000",
       "The minimum task size for sharding parallel operation.">
   ];
 
   let dependentDialects = [
     "arith::ArithmeticDialect",

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

 extern "C" void mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *, CoroHandle,
                                                      CoroResume);
 
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime after the all members of the group become ready.
 extern "C" void
 mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume);
 
+// Returns the current number of available worker threads in the threadpool.
+extern "C" int64_t mlirAsyncRuntimGetNumWorkerThreads();
 //===----------------------------------------------------------------------===//
 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//
 
 extern "C" void mlirAsyncRuntimePrintCurrentThreadId();
 
 } // namespace runtime
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

 static constexpr const char *kAddTokenToGroup =
     "mlirAsyncRuntimeAddTokenToGroup";
 static constexpr const char *kAwaitTokenAndExecute =
     "mlirAsyncRuntimeAwaitTokenAndExecute";
 static constexpr const char *kAwaitValueAndExecute =
     "mlirAsyncRuntimeAwaitValueAndExecute";
 static constexpr const char *kAwaitAllAndExecute =
     "mlirAsyncRuntimeAwaitAllInGroupAndExecute";
+static constexpr const char *kGetNumWorkerThreads =
+    "mlirAsyncRuntimGetNumWorkerThreads";
 
 namespace {
 /// Async Runtime API function types.
 ///
 /// Because we can't create API function signature for type parametrized
 /// async.value type, we use opaque pointers (!llvm.ptr<i8>) instead. After
 /// lowering all async data types become opaque pointers at runtime.
 struct AsyncAPI {
struct AsyncAPI {
   }
 
   static FunctionType awaitAllAndExecuteFunctionType(MLIRContext *ctx) {
     auto hdl = opaquePointerType(ctx);
     auto resume = LLVM::LLVMPointerType::get(resumeFunctionType(ctx));
     return FunctionType::get(ctx, {GroupType::get(ctx), hdl, resume}, {});
   }
 
+  static FunctionType getNumWorkerThreads(MLIRContext *ctx) {
+    return FunctionType::get(ctx, {}, {IndexType::get(ctx)});
+  }
+
   // Auxiliary coroutine resume intrinsic wrapper.
   static Type resumeFunctionType(MLIRContext *ctx) {
     auto voidTy = LLVM::LLVMVoidType::get(ctx);
     auto i8Ptr = opaquePointerType(ctx);
     return LLVM::LLVMFunctionType::get(voidTy, {i8Ptr}, false);
   }
 };
 } // namespace
static void addAsyncRuntimeApiDeclarations(ModuleOp module) {
   addFuncDecl(kGetValueStorage, AsyncAPI::getValueStorageFunctionType(ctx));
   addFuncDecl(kAddTokenToGroup, AsyncAPI::addTokenToGroupFunctionType(ctx));
   addFuncDecl(kAwaitTokenAndExecute,
               AsyncAPI::awaitTokenAndExecuteFunctionType(ctx));
   addFuncDecl(kAwaitValueAndExecute,
               AsyncAPI::awaitValueAndExecuteFunctionType(ctx));
   addFuncDecl(kAwaitAllAndExecute,
               AsyncAPI::awaitAllAndExecuteFunctionType(ctx));
+  addFuncDecl(kGetNumWorkerThreads, AsyncAPI::getNumWorkerThreads(ctx));
 }
 
 //===----------------------------------------------------------------------===//
 // Add malloc/free declarations to the module.
 //===----------------------------------------------------------------------===//
 
 static constexpr const char *kMalloc = "malloc";
 static constexpr const char *kFree = "free";
rewriter.replaceOpWithNewOp<CallOp>(
         op, kAddTokenToGroup, rewriter.getI64Type(), adaptor.getOperands());
 
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// Convert async.runtime.num_worker_threads to the corresponding runtime API
+// call.
+//===----------------------------------------------------------------------===//
+
+namespace {
+class RuntimeNumWorkerThreadsOpLowering
+    : public OpConversionPattern<RuntimeNumWorkerThreadsOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(RuntimeNumWorkerThreadsOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    // Replace with a runtime API function call.
+    rewriter.replaceOpWithNewOp<CallOp>(op, kGetNumWorkerThreads,
+                                        rewriter.getIndexType());
+
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
 // Async reference counting ops lowering (`async.runtime.add_ref` and
 // `async.runtime.drop_ref` to the corresponding API calls).
 //===----------------------------------------------------------------------===//
 
 namespace {
 template <typename RefCountingOp>
 class RefCountingOpLowering : public OpConversionPattern<RefCountingOp> {
 public:
void ConvertAsyncToLLVMPass::runOnOperation() {
 
   // Convert return operations inside async.execute regions.
   patterns.add<ReturnOpOpConversion>(converter, ctx);
 
   // Lower async.runtime operations to the async runtime API calls.
   patterns.add<RuntimeSetAvailableOpLowering, RuntimeSetErrorOpLowering,
                RuntimeIsErrorOpLowering, RuntimeAwaitOpLowering,
                RuntimeAwaitAndResumeOpLowering, RuntimeResumeOpLowering,
-               RuntimeAddToGroupOpLowering, RuntimeAddRefOpLowering,
-               RuntimeDropRefOpLowering>(converter, ctx);
+               RuntimeAddToGroupOpLowering, RuntimeNumWorkerThreadsOpLowering,
+               RuntimeAddRefOpLowering, RuntimeDropRefOpLowering>(converter,
+                                                                  ctx);
 
   // Lower async.runtime operations that rely on LLVM type converter to convert
   // from async value payload type to the LLVM type.
   patterns.add<RuntimeCreateOpLowering, RuntimeCreateGroupOpLowering,
                RuntimeStoreOpLowering, RuntimeLoadOpLowering>(llvmConverter,
                                                               ctx);
 
   // Lower async coroutine operations to LLVM coroutine intrinsics.

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

for (int i = op.getNumLoops() - 2; i >= 0; --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)
         numUnrollableLoops++;
     }
 
+    Value numWorkerThreadsVal;
+    if (numWorkerThreads >= 0) {
+      numWorkerThreadsVal = b.create<arith::ConstantIndexOp>(numWorkerThreads);
+    } else {
+      numWorkerThreadsVal = b.create<async::RuntimeNumWorkerThreadsOp>();
+    }
+
     // With large number of threads the value of creating many compute blocks
     // is reduced because the problem typically becomes memory bound. For small
     // number of threads it helps with stragglers.
-    float overshardingFactor = numWorkerThreads <= 4    ? 8.0
-                               : numWorkerThreads <= 8  ? 4.0
-                               : numWorkerThreads <= 16 ? 2.0
-                               : numWorkerThreads <= 32 ? 1.0
-                               : numWorkerThreads <= 64 ? 0.8
-                                                        : 0.6;
-
-    // Do not overload worker threads with too many compute blocks.
-    Value maxComputeBlocks = b.create<arith::ConstantIndexOp>(
-        std::max(1, static_cast<int>(numWorkerThreads * overshardingFactor)));
+
+    // Pairs of upper bound of the bracket and factor that the number of workers
+    // needs to be scaled with if it falls in that bucket.
+    SmallVector<std::pair<int, float>> overshardingBrackets = {
+        {4, 8.0f}, {8, 4.0f}, {16, 2.0f}, {32, 1.0f}, {64, 0.8f}};
+
+    auto makeFloat = [](ImplicitLocOpBuilder ib, float val) -> Value {
+      return ib.create<arith::ConstantFloatOp>(llvm::APFloat(val),
+                                               ib.getF32Type());
+    };
+    std::function<Value(ImplicitLocOpBuilder &)> makeSwitch =
+        [&makeFloat](ImplicitLocOpBuilder nb) -> Value {
+      return makeFloat(nb, 0.6f);
+    };
+
+    for (std::pair<int, float> &p : llvm::reverse(overshardingBrackets)) {
+      makeSwitch = [&makeFloat, &numWorkerThreadsVal, &p,
+                    makeSwitch](ImplicitLocOpBuilder &nb) -> Value {
+        auto inBracket = nb.create<arith::CmpIOp>(
+            arith::CmpIPredicate::sle, numWorkerThreadsVal,
+            nb.create<arith::ConstantIndexOp>(p.first));
+        return nb
+            .create<scf::IfOp>(
+                TypeRange(nb.getF32Type()), inBracket,
+                [&](OpBuilder &nestedBuilder, Location nestedLoc) {
+                  ImplicitLocOpBuilder nb(nestedLoc, nestedBuilder);
+                  nb.create<scf::YieldOp>(makeFloat(nb, p.second));
+                },
+                [&](OpBuilder &nestedBuilder, Location nestedLoc) {
+                  ImplicitLocOpBuilder nb(nestedLoc, nestedBuilder);
+                  nb.create<scf::YieldOp>(makeSwitch(nb));
+                })
+            .getResults()
+            .front();
+      };
+    }
+    Value scaledWorkers = b.create<arith::IndexCastOp>(
+        b.create<arith::FPToSIOp>(
+            b.create<arith::MulFOp>(
+                makeSwitch(b), b.create<arith::SIToFPOp>(
+                                   b.create<arith::IndexCastOp>(
+                                       numWorkerThreadsVal, b.getI32Type()),
+                                   b.getF32Type())),
+            b.getI32Type()),
+        b.getIndexType());
+
+    Value maxComputeBlocks = b.create<arith::MaxSIOp>(
+        b.create<arith::ConstantIndexOp>(1), scaledWorkers);
 
     // Compute parallel block size from the parallel problem size:
     //   blockSize = min(tripCount,
     //                   max(ceil_div(tripCount, maxComputeBlocks),
     //                       minTaskSize))
     Value bs0 = b.create<arith::CeilDivSIOp>(tripCount, maxComputeBlocks);
     Value bs1 = b.create<arith::MaxSIOp>(bs0, minTaskSize);
     Value blockSize = b.create<arith::MinSIOp>(tripCount, bs1);

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

extern "C" void mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *group,
   if (group->pendingTokens == 0) {
     lock.unlock();
     execute();
   } else {
     group->awaiters.emplace_back([execute]() { execute(); });
   }
 }
 
+extern "C" int64_t mlirAsyncRuntimGetNumWorkerThreads() { return 1; }
+
 //===----------------------------------------------------------------------===//
 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//
 
 extern "C" void mlirAsyncRuntimePrintCurrentThreadId() {
   static thread_local std::thread::id thisId = std::this_thread::get_id();
   std::cout << "Current thread id: " << thisId << std::endl;
 }
void __mlir_runner_init(llvm::StringMap<void *> &exportSymbols) {
   exportSymbol("mlirAsyncRuntimeCreateGroup",
                &mlir::runtime::mlirAsyncRuntimeCreateGroup);
   exportSymbol("mlirAsyncRuntimeAddTokenToGroup",
                &mlir::runtime::mlirAsyncRuntimeAddTokenToGroup);
   exportSymbol("mlirAsyncRuntimeAwaitAllInGroup",
                &mlir::runtime::mlirAsyncRuntimeAwaitAllInGroup);
   exportSymbol("mlirAsyncRuntimeAwaitAllInGroupAndExecute",
                &mlir::runtime::mlirAsyncRuntimeAwaitAllInGroupAndExecute);
+  exportSymbol("mlirAsyncRuntimGetNumWorkerThreads",
+               &mlir::runtime::mlirAsyncRuntimGetNumWorkerThreads);
   exportSymbol("mlirAsyncRuntimePrintCurrentThreadId",
                &mlir::runtime::mlirAsyncRuntimePrintCurrentThreadId);
 }
 
 // NOLINTNEXTLINE(*-identifier-naming): externally called.
 extern "C" API void __mlir_runner_destroy() { resetDefaultAsyncRuntime(); }
 
 } // namespace runtime
 } // namespace mlir
 
 #endif // MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS

mlir/test/Dialect/Async/async-parallel-for-num-worker-threads.mlir

+// RUN: mlir-opt %s -split-input-file -async-parallel-for=num-workers=-1  \
+// RUN: | FileCheck %s --dump-input=always
+
+// We check only a single level of recursion in the switch statement in order
+// to avoid brittleness.
+
+// CHECK-LABEL: @num_worker_threads(
+// CHECK:       %[[MEMREF:.*]]: memref<?xf32>
+func @num_worker_threads(%arg0: memref<?xf32>) {
+
+  // CHECK: %[[scalingCst:.*]] = arith.constant {{.*}} : f32
+  // CHECK: scf.if %false {
+  // CHECK: } else {
+  // CHECK:   %[[workersIndex:.*]] = async.runtime.num_worker_threads
+  // CHECK:   %[[cond:.*]] = arith.cmpi sle, %[[workersIndex]], %c4 : index
+  // CHECK:   %[[factor:.*]] = scf.if %[[cond]] -> (f32) {
+  // CHECK:     scf.yield {{.*}} : f32
+  // CHECK:   } else {
+  // Skipping redundancy...
+  // CHECK:   }
+  // CHECK:   %[[workersInt:.*]] = arith.index_cast %[[workersIndex]] : index to i32
+  // CHECK:   %[[workersFloat:.*]] = arith.sitofp %[[workersInt]] : i32 to f32
+  // CHECK:   %[[scaledFloat:.*]] = arith.mulf %[[factor]], %[[workersFloat]] : f32
+  // CHECK:   %[[scaledInt:.*]] = arith.fptosi %[[scaledFloat]] : f32 to i32
+  // CHECK:   %[[scaledIndex:.*]] = arith.index_cast %[[scaledInt]] : i32 to index
+
+  %lb = arith.constant 0 : index
+  %ub = arith.constant 100 : index
+  %st = arith.constant 1 : index
+  scf.parallel (%i) = (%lb) to (%ub) step (%st) {
+    %one = arith.constant 1.0 : f32
+    memref.store %one, %arg0[%i] : memref<?xf32>
+  }
+
+  return
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

td_library(
     srcs = [
         "include/mlir/Dialect/Async/IR/AsyncDialect.td",
         "include/mlir/Dialect/Async/IR/AsyncOps.td",
         "include/mlir/Dialect/Async/IR/AsyncTypes.td",
     ],
     includes = ["include"],
     deps = [
         ":ControlFlowInterfacesTdFiles",
+        ":InferTypeOpInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "AsyncOpsIncGen",
     strip_include_prefix = "include",
cc_library(
     ]),
     hdrs = glob([
         "include/mlir/Dialect/Async/IR/*.h",
     ]),
     includes = ["include"],
     deps = [
         ":AsyncOpsIncGen",
         ":ControlFlowInterfaces",
-        ":Dialect",
         ":IR",
-        ":SideEffectInterfaces",
-        ":StandardOps",
-        ":Support",
+        ":InferTypeOpInterface",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "AsyncTransforms",
     srcs = glob([
         "lib/Dialect/Async/Transforms/*.cpp",