[OSPP][Schedule] Design and Implementation of Scheduling Module

examples/ScheDialect/.gitignore

@@ -0,0 +1,3 @@
+log.mlir
+log.ll
+log.s

examples/ScheDialect/makefile

@@ -0,0 +1,58 @@
+#!/bin/bash
+BUDDY_OPT := ../../build/bin/buddy-opt
+MLIR_OPT := ../../llvm/build/bin/mlir-opt
+MLIR_TRANSLATE := ../../llvm/build/bin/mlir-translate
+MLIR_CPU_RUNNER := ../../llvm/build/bin/mlir-cpu-runner
+LLC := ../../llvm/build/bin/llc
+OPT_FLAG := -O0
+
+ifeq ($(shell uname),Linux)
+MLIR_RUNNER_UTILS := ../../llvm/build/lib/libmlir_runner_utils.so
+MLIR_C_RUNNER_UTILS := ../../llvm/build/lib/libmlir_c_runner_utils.so
+MLIR_ASYNC_RUNTIME := ../../llvm/build/lib/libmlir_async_runtime.so
+MLIR_CUDA_RUNTIME := ../../llvm/build/lib/libmlir_cuda_runtime.so
+MTRIPLE := x86_64-unknown-linux-gnu
+else ifeq ($(shell uname),Darwin)
+MLIR_RUNNER_UTILS := ../../llvm/build/lib/libmlir_runner_utils.dylib
+MLIR_C_RUNNER_UTILS := ../../llvm/build/lib/libmlir_c_runner_utils.dylib
+MLIR_ASYNC_RUNTIME := ./../llvm/build/lib/libmlir_async_runtime.dylib
+MTRIPLE := x86_64-apple-darwin
+endif
+
+.SECONDEXPANSION:
+all-run: $$(run-targets)
+
+sche-addf-tensor-lower:
+	@${BUDDY_OPT} ./sche-addf-tensor.mlir \
+		--device-schedule \
+		--lower-sche -o ./log.mlir
+
+sche-addf-vector-lower:
+	@${BUDDY_OPT} ./sche-addf-vector.mlir \
+		--device-schedule \
+		--lower-sche -o ./log.mlir
+
+sche-scf-for-lower:
+	@${BUDDY_OPT} ./sche-scf-for.mlir -device-schedule | \
+	${BUDDY_OPT} -lower-sche | \
+	${MLIR_OPT} -gpu-kernel-outlining | \
+	${MLIR_OPT} --pass-pipeline="builtin.module(func.func(tosa-to-linalg),empty-tensor-to-alloc-tensor,arith-bufferize, func.func(linalg-bufferize), func.func(tensor-bufferize), func-bufferize,func.func(buffer-deallocation, convert-linalg-to-loops) , convert-scf-to-cf, expand-strided-metadata, memref-expand)" | \
+	${MLIR_OPT} -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,gpu-to-cubin))' | \
+	${MLIR_OPT} -gpu-async-region -gpu-to-llvm -o log.mlir
+
+run-targets += sche-scf-for-run
+sche-scf-for-run:
+	@${BUDDY_OPT} ./sche-scf-for.mlir -device-schedule | \
+	${BUDDY_OPT} -lower-sche | \
+	${MLIR_OPT} -gpu-kernel-outlining | \
+	${MLIR_OPT} --pass-pipeline="builtin.module(func.func(tosa-to-linalg),empty-tensor-to-alloc-tensor,arith-bufferize, func.func(linalg-bufferize), func.func(tensor-bufferize), func-bufferize,func.func(buffer-deallocation, convert-linalg-to-loops) , convert-scf-to-cf, expand-strided-metadata, memref-expand)" | \
+	${MLIR_OPT} -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,gpu-to-cubin))' | \
+	${MLIR_OPT} -gpu-async-region -gpu-to-llvm | \
+	${MLIR_OPT} -async-to-async-runtime -async-runtime-ref-counting | \
+	${MLIR_OPT} -convert-async-to-llvm -convert-func-to-llvm | \
+	${MLIR_CPU_RUNNER} -entry-point-result=void -shared-libs=${MLIR_RUNNER_UTILS} -shared-libs=${MLIR_CUDA_RUNTIME} \
+	-shared-libs=${MLIR_ASYNC_RUNTIME} ${OPT_FLAG}
+
+sche-reduce-add-lower:
+	@${BUDDY_OPT} ./sche-reduce-add.mlir \
+		--device-schedule -o ./log.mlir

examples/ScheDialect/sche-scf-for.mlir

@@ -0,0 +1,50 @@
+// RUN: buddy-opt %s -device-schedule \
+// RUN: -device-schedule -lower-sche \
+// RUN:	-gpu-kernel-outlining \
+// RUN:	| buddy-opt --pass-pipeline="builtin.module(func.func(tosa-to-linalg),empty-tensor-to-alloc-tensor,arith-bufferize, func.func(linalg-bufferize), func.func(tensor-bufferize), func-bufferize,func.func(buffer-deallocation, convert-linalg-to-loops) , convert-scf-to-cf, expand-strided-metadata, memref-expand)" \
+// RUN:	| buddy-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,gpu-to-cubin))' \
+// RUN:	| buddy-opt -gpu-async-region -gpu-to-llvm \
+// RUN:	-async-to-async-runtime -async-runtime-ref-counting \
+// RUN:	-convert-async-to-llvm -convert-func-to-llvm \
+// RUN: | mlir-cpu-runner  -entry-point-result=void \
+// RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_cuda_runtime%shlibext \
+// RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext \
+// RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_async_runtime%shlibext \
+// RUN: | FileCheck %s
+
+func.func @main() {
+  %c0 = arith.constant 0.0 : f32
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c8 = arith.constant 8 : index
+
+  %lb = arith.constant 0 : index
+  %ub = arith.constant 100 : index
+
+  %A = memref.alloc() : memref<100xf32>
+  %A_cast0 = memref.cast %A : memref<100xf32> to memref<*xf32>
+
+  scf.for %iv = %lb to %ub step %c1 {
+    memref.store %c0, %A[%iv] : memref<100xf32>
+  }
+
+  scf.for %iv = %lb to %ub step %c1 {
+    %0 = arith.muli %iv, %c8 : index
+    %1 = arith.addi %iv, %0  : index
+    %2 = arith.index_cast %1 : index to i32
+    %3 = arith.sitofp %2 : i32 to f32
+    %4 = memref.load %A[%iv] : memref<100xf32>
+    %5 = arith.addf %3, %4 : f32
+    memref.store %5, %A[%iv] : memref<100xf32>
+  } {sche.devices = [{targetId = "cpu", targetConfig = "", duty_ratio = 0.2:f32}, {targetId = "gpu", targetConfig = "", duty_ratio = 0.8:f32}]}
+
+  %res = memref.load %A[%c1] : memref<100xf32>
+  // CHECK: 10
+  call @printMemrefF32(%A_cast0) : (memref<*xf32>) -> ()
+
+  memref.dealloc %A : memref<100xf32>
+
+  return
+}
+
+func.func private @printMemrefF32(%ptr : memref<*xf32>)

examples/lit.cfg.py

@@ -54,6 +54,7 @@
     'RISCVBuddyExt',
     'RVVDialect',
     'RVVExperiment',
+    'ScheDialect',
     'SIMDExperiment',
     'ToyDSL',
     'VectorExpDialect',

midend/include/Dialect/CMakeLists.txt

@@ -4,3 +4,4 @@ add_subdirectory(DIP)
 add_subdirectory(RVV)
 add_subdirectory(VectorExp)
 add_subdirectory(Gemmini)
+add_subdirectory(Sche)

midend/include/Dialect/Sche/CMakeLists.txt

@@ -0,0 +1,16 @@
+set(LLVM_TARGET_DEFINITIONS ScheOps.td)
+mlir_tablegen(ScheOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(ScheOpsEnums.cpp.inc -gen-enum-defs)
+
+set(LLVM_TARGET_DEFINITIONS ScheOps.td)
+mlir_tablegen(ScheOpsAttributes.h.inc -gen-attrdef-decls)
+mlir_tablegen(ScheOpsAttributes.cpp.inc -gen-attrdef-defs)
+
+set(LLVM_TARGET_DEFINITIONS ScheDialect.td)
+mlir_tablegen(ScheOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(ScheOpInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRScheOpInterfacesIncGen)
+
+add_mlir_dialect(ScheOps sche)
+# add_mlir_doc(ScheDialect -gen-dialect-doc ScheDialect Sche/)
+# add_mlir_doc(ScheOps -gen-op-doc ScheOps Sche/)

midend/include/Dialect/Sche/ScheDialect.h

@@ -0,0 +1,43 @@
+//===- ScheDialect.h - Sche Dialect Definition --------------------*- C++ -*-===//
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+// 
+// This is the header file for the sche dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCHE_SCHEDIALECT_H
+#define SCHE_SCHEDIALECT_H
+
+namespace buddy {
+namespace sche {
+class AsyncTokenType
+    : public ::mlir::Type::TypeBase<AsyncTokenType, ::mlir::Type, ::mlir::TypeStorage> {
+public:
+  // Used for generic hooks in TypeBase.
+  using Base::Base;
+};
+
+void addAsyncDependency(::mlir::Operation *op, ::mlir::Value token);
+}
+}
+
+#include "mlir/IR/Dialect.h"
+
+#include "Sche/ScheOpsDialect.h.inc"
+
+#include "Sche/ScheOpInterfaces.h.inc"
+
+#endif // SCHE_SCHEDIALECT_H

midend/include/Dialect/Sche/ScheDialect.td

@@ -0,0 +1,103 @@
+//===- ScheDialect.td - Sche Dialect Definition --------------*- tablegen -*-===//
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the top level file for the sche dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCHE_SCHEDIALECT_TD
+#define SCHE_SCHEDIALECT_TD
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// Sche Dialect Definition.
+//===----------------------------------------------------------------------===//
+
+def Sche_Dialect : Dialect {
+  let name = "sche";
+  let summary = "The Sche Dialect.";
+  let description = [{
+    The `sche` dialect is for scheduling operators on heterogeneous devices.
+  }];
+  let cppNamespace = "::buddy::sche";
+  // let useDefaultAttributePrinterParser = 1;
+  let useDefaultTypePrinterParser = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Sche Interfaces.
+//===----------------------------------------------------------------------===//
+
+def Sche_AsyncOpInterface : OpInterface<"AsyncOpInterface"> {
+  let description = [{
+    Interface for GPU operations that execute asynchronously on the device.
+
+    GPU operations implementing this interface take a list of dependencies
+    as `gpu.async.token` arguments and optionally return a `gpu.async.token`.
+
+    The op doesn't start executing until all depent ops producing the async
+    dependency tokens have finished executing.
+
+    If the op returns a token, the op merely schedules the execution on the
+    device and returns immediately, without waiting for the execution to
+    complete. On the hand, if the op does not return a token, the op will wait
+    for the execution to complete.
+  }];
+  let cppNamespace = "::buddy::sche";
+
+  let methods = [
+    InterfaceMethod<[{
+        Query the operands that represent async dependency tokens.
+      }],
+      "::mlir::OperandRange", "getAsyncDependencies", (ins), [{}], [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getAsyncDependencies();
+      }]
+    >,
+    InterfaceMethod<[{
+        Adds a new token to the list of async dependencies if it is not already there.
+      }],
+      "void", "addAsyncDependency", (ins "::mlir::Value":$token),
+      [{}], [{
+        if (!::llvm::is_contained(this->getAsyncDependencies(), token))
+          ::buddy::sche::addAsyncDependency(this->getOperation(), token);
+      }]
+    >,
+    InterfaceMethod<[{
+        Query the result that represents the async token to depend on.
+      }],
+      "::mlir::Value", "getAsyncToken"
+    >
+  ];
+}
+
+//===----------------------------------------------------------------------===//
+// Sche Types.
+//===----------------------------------------------------------------------===//
+
+def Sche_AsyncToken : DialectType<
+  Sche_Dialect, CPred<"$_self.isa<::buddy::sche::AsyncTokenType>()">, "async token type">,
+             BuildableType<"buddy::sche::AsyncTokenType::get($_builder.getContext())">;
+
+//===----------------------------------------------------------------------===//
+// Base Sche Operation Definition.
+//===----------------------------------------------------------------------===//
+
+class Sche_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Sche_Dialect, mnemonic, traits>;
+
+#endif // SCHE_SCHEDIALECT_TD

midend/include/Dialect/Sche/ScheOps.h

@@ -0,0 +1,41 @@
+//===- ScheOps.h - Sche Dialect Ops -------------------------------*- C++ -*-===//
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the header file for operations in sche dialect. 
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCHE_SCHEOPS_H
+#define SCHE_SCHEOPS_H
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "Dialect/Sche/ScheDialect.h"
+#include "mlir/IR/OpImplementation.h"
+
+// #include "Sche/ScheOpsEnums.h.inc"
+
+// #define GET_ATTRDEF_CLASSES
+// #include "Sche/ScheOpsAttributes.h.inc"
+
+#define GET_OP_CLASSES
+#include "Sche/ScheOps.h.inc"
+
+#endif // SCHE_SCHEOPS_H