better throughput scheduling with steady state cost model, now limited to spatial
This commit is contained in:
NiccoloN
@@ -0,0 +1,158 @@
|
||||
#include "TestCheck.hpp"
|
||||
#include <cmath>
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/IR/BuiltinOps.h"
|
||||
#include "mlir/IR/Verifier.h"
|
||||
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
using namespace onnx_mlir;
|
||||
using namespace onnx_mlir::spatial;
|
||||
|
||||
static void testZeroContent(OpBuilder &builder) {
|
||||
auto type = RankedTensorType::get({4}, builder.getF32Type());
|
||||
for (double value : {0.0, -0.0, 1.0, -1.0}) {
|
||||
auto attribute = DenseElementsAttr::get(type, builder.getF32FloatAttr(value));
|
||||
TEST_CHECK(isZeroBitPatternSplat(attribute) == (value == 0.0 && !std::signbit(value)));
|
||||
}
|
||||
auto integerType = RankedTensorType::get({4}, builder.getI32Type());
|
||||
TEST_CHECK(isZeroBitPatternSplat(DenseElementsAttr::get(integerType, builder.getI32IntegerAttr(0))));
|
||||
TEST_CHECK(!isZeroBitPatternSplat(DenseElementsAttr::get(integerType, builder.getI32IntegerAttr(1))));
|
||||
TEST_CHECK(!isZeroBitPatternSplat(DenseElementsAttr::get(type, ArrayRef<float>{0, 1, 0, 0})));
|
||||
TEST_CHECK(!isZeroBitPatternSplat(builder.getF32FloatAttr(0)));
|
||||
TEST_CHECK(!isZeroBitPatternSplat({}));
|
||||
}
|
||||
|
||||
static void testInputPlan(MLIRContext &context, bool batched) {
|
||||
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
|
||||
OpBuilder builder(module->getBodyRegion());
|
||||
auto loc = module->getLoc();
|
||||
auto type = RankedTensorType::get({1, 4}, builder.getF32Type());
|
||||
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({type}, {}));
|
||||
builder.setInsertionPointToStart(function.addEntryBlock());
|
||||
Value constant = arith::ConstantOp::create(builder, loc,
|
||||
DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
|
||||
Value view = tensor::CastOp::create(builder, loc, type, constant);
|
||||
SmallVector<Value> inputs{constant, function.getArgument(0), constant, view, function.getArgument(0)};
|
||||
Operation *owner;
|
||||
if (batched) {
|
||||
auto batch = SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
|
||||
builder.getI32IntegerAttr(2), ValueRange{}, inputs);
|
||||
batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({0, 1}));
|
||||
owner = batch;
|
||||
} else {
|
||||
auto single = SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, inputs);
|
||||
single->setAttr("coreId", builder.getI32IntegerAttr(0));
|
||||
owner = single;
|
||||
}
|
||||
auto *body = new Block;
|
||||
owner->getRegion(0).push_back(body);
|
||||
if (batched) body->addArgument(builder.getIndexType(), loc);
|
||||
for (Value input : inputs) body->addArgument(input.getType(), loc);
|
||||
builder.setInsertionPointToStart(body);
|
||||
unsigned first = batched ? 1 : 0;
|
||||
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
|
||||
Value limit = arith::ConstantIndexOp::create(builder, loc, 1000000000000);
|
||||
Value one = arith::ConstantIndexOp::create(builder, loc, 1);
|
||||
auto loop = scf::ForOp::create(builder, loc, zero, limit, one);
|
||||
builder.setInsertionPointToStart(loop.getBody());
|
||||
SpatVAddOp::create(builder, loc, type, body->getArgument(first), body->getArgument(first + 1));
|
||||
SpatVAddOp::create(builder, loc, type, body->getArgument(first + 3), body->getArgument(first));
|
||||
builder.setInsertionPointAfter(loop);
|
||||
auto lastUse = SpatVAddOp::create(builder, loc, type, body->getArgument(first + 2), body->getArgument(first));
|
||||
auto lastFunctionUse = SpatVAddOp::create(builder, loc, type,
|
||||
body->getArgument(first + 4), body->getArgument(first));
|
||||
SpatYieldOp::create(builder, loc, ValueRange{});
|
||||
builder.setInsertionPointAfter(owner);
|
||||
func::ReturnOp::create(builder, loc);
|
||||
TEST_CHECK(succeeded(verify(*module)));
|
||||
// Body lowering can remove an input's final use before input materialization.
|
||||
lastUse.erase();
|
||||
lastFunctionUse.erase();
|
||||
|
||||
std::string before;
|
||||
llvm::raw_string_ostream stream(before);
|
||||
module->print(stream);
|
||||
for (unsigned repeat = 0; repeat < 2; ++repeat) {
|
||||
auto constantPlan = planComputeInputMaterialization(owner, 0);
|
||||
auto runtimePlan = planComputeInputMaterialization(owner, 1);
|
||||
auto unusedPlan = planComputeInputMaterialization(owner, 2);
|
||||
auto viewPlan = planComputeInputMaterialization(owner, 3);
|
||||
TEST_CHECK(succeeded(constantPlan) && succeeded(runtimePlan) && succeeded(unusedPlan));
|
||||
TEST_CHECK(*constantPlan == (batched ? ComputeInputMaterialization::CopyAtEntry
|
||||
: ComputeInputMaterialization::ForwardConstant));
|
||||
TEST_CHECK(*runtimePlan == ComputeInputMaterialization::ForwardFunctionInput);
|
||||
auto unusedFunctionPlan = planComputeInputMaterialization(owner, 4);
|
||||
TEST_CHECK(succeeded(unusedFunctionPlan)
|
||||
&& *unusedFunctionPlan == (batched ? ComputeInputMaterialization::ForwardFunctionInput
|
||||
: ComputeInputMaterialization::Unused));
|
||||
TEST_CHECK(*unusedPlan == (batched ? ComputeInputMaterialization::CopyAtEntry
|
||||
: ComputeInputMaterialization::Unused));
|
||||
TEST_CHECK(succeeded(viewPlan) && *viewPlan == ComputeInputMaterialization::CopyAtEntry);
|
||||
}
|
||||
TEST_CHECK(failed(planComputeInputMaterialization(owner, 5)));
|
||||
TEST_CHECK(failed(planComputeInputMaterialization(function, 0)));
|
||||
std::string after;
|
||||
llvm::raw_string_ostream afterStream(after);
|
||||
module->print(afterStream);
|
||||
TEST_CHECK(before == after);
|
||||
}
|
||||
|
||||
static void testOutputReuse(MLIRContext &context) {
|
||||
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
|
||||
IRRewriter builder(&context);
|
||||
builder.setInsertionPointToStart(module->getBody());
|
||||
auto loc = module->getLoc();
|
||||
auto type = RankedTensorType::get({1, 4}, builder.getF32Type());
|
||||
auto function = func::FuncOp::create(builder, loc, "reuse", builder.getFunctionType({}, {}));
|
||||
builder.setInsertionPointToStart(function.addEntryBlock());
|
||||
Value input = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
|
||||
Value constant = arith::ConstantOp::create(builder, loc,
|
||||
DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
|
||||
auto first = SpatVAddOp::create(builder, loc, type, input, constant);
|
||||
TEST_CHECK(getReusableOutputTensor(first) == input);
|
||||
auto second = SpatVAddOp::create(builder, loc, type, input, constant);
|
||||
TEST_CHECK(!getReusableOutputTensor(first));
|
||||
TEST_CHECK(getReusableOutputTensor(second) == input);
|
||||
second->moveBefore(first);
|
||||
TEST_CHECK(getReusableOutputTensor(first) == input);
|
||||
TEST_CHECK(!getReusableOutputTensor(second));
|
||||
second->moveAfter(first);
|
||||
TEST_CHECK(!getReusableOutputTensor(first));
|
||||
TEST_CHECK(getReusableOutputTensor(second) == input);
|
||||
Value view = tensor::CastOp::create(builder, loc, type, input);
|
||||
auto fromView = SpatVAddOp::create(builder, loc, type, view, constant);
|
||||
TEST_CHECK(!getReusableOutputTensor(fromView));
|
||||
TEST_CHECK(getBestOutputTensorFromOperandsOrAllocate(builder, fromView)
|
||||
.getDefiningOp<tensor::EmptyOp>());
|
||||
builder.setInsertionPointAfter(fromView);
|
||||
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
|
||||
Value limit = arith::ConstantIndexOp::create(builder, loc, 10);
|
||||
Value one = arith::ConstantIndexOp::create(builder, loc, 1);
|
||||
auto loop = scf::ForOp::create(builder, loc, zero, limit, one, ValueRange{input});
|
||||
builder.setInsertionPointToStart(loop.getBody());
|
||||
auto carried = SpatVAddOp::create(builder, loc, type, loop.getRegionIterArg(0), constant);
|
||||
scf::YieldOp::create(builder, loc, ValueRange{carried});
|
||||
TEST_CHECK(getReusableOutputTensor(carried) == loop.getRegionIterArg(0));
|
||||
TEST_CHECK(getBestOutputTensorFromOperandsOrAllocate(builder, carried) == loop.getRegionIterArg(0));
|
||||
builder.setInsertionPointAfter(loop);
|
||||
func::ReturnOp::create(builder, loc);
|
||||
TEST_CHECK(succeeded(verify(*module)));
|
||||
}
|
||||
|
||||
int main() {
|
||||
MLIRContext context(MLIRContext::Threading::DISABLED);
|
||||
context.loadDialect<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
|
||||
tensor::TensorDialect, SpatialDialect>();
|
||||
OpBuilder builder(&context);
|
||||
testZeroContent(builder);
|
||||
testInputPlan(context, false);
|
||||
testInputPlan(context, true);
|
||||
testOutputReuse(context);
|
||||
}
|
||||
Reference in New Issue
Block a user