[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

800

Nit: drop trivial braces.

818

Please avoid magic numbers.

821

const &

829

TypeRange should be implicitly constructible from a single type.

838–839

.getResult(0) looks shorter.

842–853

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

13

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

16

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

129

nit: empty line before the next comment

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

800

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

809

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

809

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

[ezhulenev]

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

808

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

817

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

821

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

800

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.

808

Done.

809

I'm assuming this is withdrawn.

817

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.

821

Great!

829

Obsolete.

838–839

Obsolete.

842–853

Done.

[ezhulenev]

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

822

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",
+           [DeclareOpInterfaceMethods<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 `:` type($result)";
+}
+
 #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/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRAsync
   Async.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Async
 
   DEPENDS
   MLIRAsyncOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
+  MLIRInferTypeOpInterface
   MLIRIR
   )

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

 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/Async/Transforms.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/RegionUtils.h"
 
 using namespace mlir;
 using namespace mlir::async;
 
 #define DEBUG_TYPE "async-parallel-for"
 
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)));
+    // is reduced because the problem typically becomes memory bound. For this
+    // reason we scale the number of workers using an equivalent to the
+    // following logic:
+    //   float overshardingFactor = numWorkerThreads <= 4    ? 8.0
+    //                              : numWorkerThreads <= 8  ? 4.0
+    //                              : numWorkerThreads <= 16 ? 2.0
+    //                              : numWorkerThreads <= 32 ? 1.0
+    //                              : numWorkerThreads <= 64 ? 0.8
+    //                                                       : 0.6;
+
+    // Pairs of non-inclusive lower end of the bracket and factor that the
+    // number of workers needs to be scaled with if it falls in that bucket.
+    const SmallVector<std::pair<int, float>> overshardingBrackets = {
+        {4, 4.0f}, {8, 2.0f}, {16, 1.0f}, {32, 0.8f}, {64, 0.6f}};
+    const float initialOvershardingFactor = 8.0f;
+
+    Value scalingFactor = b.create<arith::ConstantFloatOp>(
+        llvm::APFloat(initialOvershardingFactor), b.getF32Type());
+    for (const std::pair<int, float> &p : overshardingBrackets) {
+      Value bracketBegin = b.create<arith::ConstantIndexOp>(p.first);
+      Value inBracket = b.create<arith::CmpIOp>(
+          arith::CmpIPredicate::sgt, numWorkerThreadsVal, bracketBegin);
+      Value bracketScalingFactor = b.create<arith::ConstantFloatOp>(
+          llvm::APFloat(p.second), b.getF32Type());
+      scalingFactor =
+          b.create<SelectOp>(inBracket, bracketScalingFactor, scalingFactor);
+    }
+    Value numWorkersIndex =
+        b.create<arith::IndexCastOp>(numWorkerThreadsVal, b.getI32Type());
+    Value numWorkersFloat =
+        b.create<arith::SIToFPOp>(numWorkersIndex, b.getF32Type());
+    Value scaledNumWorkers =
+        b.create<arith::MulFOp>(scalingFactor, numWorkersFloat);
+    Value scaledNumInt =
+        b.create<arith::FPToSIOp>(scaledNumWorkers, b.getI32Type());
+    Value scaledWorkers =
+        b.create<arith::IndexCastOp>(scaledNumInt, 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);
 
     ParallelComputeFunction notUnrollableParallelComputeFunction =
         createParallelComputeFunction(op, staticBounds, 0, rewriter);
 
     // Dispatch parallel compute function using async recursive work splitting,
     // or by submitting compute task sequentially from a caller thread.
     auto doDispatch = asyncDispatch ? doAsyncDispatch : doSequentialDispatch;
 

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 getDefaultAsyncRuntime()->getThreadPool().getThreadCount();
+}
+
 //===----------------------------------------------------------------------===//
 // 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
+
+// CHECK-LABEL: @num_worker_threads(
+// CHECK:       %[[MEMREF:.*]]: memref<?xf32>
+func @num_worker_threads(%arg0: memref<?xf32>) {
+
+  // CHECK: %[[scalingCstInit:.*]] = arith.constant 8.000000e+00 : f32
+  // CHECK: %[[bracketLowerBound4:.*]] = arith.constant 4 : index
+  // CHECK: %[[scalingCst4:.*]] = arith.constant 4.000000e+00 : f32
+  // CHECK: %[[bracketLowerBound8:.*]] = arith.constant 8 : index
+  // CHECK: %[[scalingCst8:.*]] = arith.constant 2.000000e+00 : f32
+  // CHECK: %[[bracketLowerBound16:.*]] = arith.constant 16 : index
+  // CHECK: %[[scalingCst16:.*]] = arith.constant 1.000000e+00 : f32
+  // CHECK: %[[bracketLowerBound32:.*]] = arith.constant 32 : index
+  // CHECK: %[[scalingCst32:.*]] = arith.constant 8.000000e-01 : f32
+  // CHECK: %[[bracketLowerBound64:.*]] = arith.constant 64 : index
+  // CHECK: %[[scalingCst64:.*]] = arith.constant 6.000000e-01 : f32
+  // CHECK:   %[[workersIndex:.*]] = async.runtime.num_worker_threads : index
+  // CHECK:   %[[inBracket4:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound4]] : index
+  // CHECK:   %[[scalingFactor4:.*]] = select %[[inBracket4]], %[[scalingCst4]], %[[scalingCstInit]] : f32
+  // CHECK:   %[[inBracket8:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound8]] : index
+  // CHECK:   %[[scalingFactor8:.*]] = select %[[inBracket8]], %[[scalingCst8]], %[[scalingFactor4]] : f32
+  // CHECK:   %[[inBracket16:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound16]] : index
+  // CHECK:   %[[scalingFactor16:.*]] = select %[[inBracket16]], %[[scalingCst16]], %[[scalingFactor8]] : f32
+  // CHECK:   %[[inBracket32:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound32]] : index
+  // CHECK:   %[[scalingFactor32:.*]] = select %[[inBracket32]], %[[scalingCst32]], %[[scalingFactor16]] : f32
+  // CHECK:   %[[inBracket64:.*]] = arith.cmpi sgt, %[[workersIndex]], %[[bracketLowerBound64]] : index
+  // CHECK:   %[[scalingFactor64:.*]] = select %[[inBracket64]], %[[scalingCst64]], %[[scalingFactor32]] : f32
+  // CHECK:   %[[workersInt:.*]] = arith.index_cast %[[workersIndex]] : index to i32
+  // CHECK:   %[[workersFloat:.*]] = arith.sitofp %[[workersInt]] : i32 to f32
+  // CHECK:   %[[scaledFloat:.*]] = arith.mulf %[[scalingFactor64]], %[[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",
hdrs = glob([
         "include/mlir/Dialect/Async/IR/*.h",
     ]),
     includes = ["include"],
     deps = [
         ":AsyncOpsIncGen",
         ":ControlFlowInterfaces",
         ":Dialect",
         ":IR",
+        ":InferTypeOpInterface",
         ":SideEffectInterfaces",
         ":StandardOps",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 cc_library(