Split toy dialect using static registration

1 files changed, 0 insertions, 452 deletions

diff --git a/mlir/LateLowering.cpp b/mlir/LateLowering.cpp
deleted file mode 100644
index eeae6ee..0000000
--- a/mlir/LateLowering.cpp
+++ /dev/null
@@ -1,452 +0,0 @@
-//====- LateLowering.cpp - Lowering from Toy+Linalg to LLVM -===//
-//
-// Copyright 2019 The MLIR Authors.
-//
-// 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 file implements late lowering of IR mixing Toy and Linalg to LLVM.
-// It involves intemerdiate steps:
-// -
-// - a mix of affine and standard dialect.
-//
-//===----------------------------------------------------------------------===//
-
-#include "toy/Dialect.h"
-
-#include "linalg3/Intrinsics.h"
-#include "linalg1/ViewOp.h"
-#include "linalg3/ConvertToLLVMDialect.h"
-#include "linalg3/TensorOps.h"
-#include "linalg3/Transforms.h"
-#include "mlir/EDSC/Builders.h"
-#include "mlir/EDSC/Helpers.h"
-#include "mlir/EDSC/Intrinsics.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/OperationSupport.h"
-#include "mlir/IR/StandardTypes.h"
-#include "mlir/LLVMIR/LLVMDialect.h"
-#include "mlir/Parser.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/DialectConversion.h"
-
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/IR/DerivedTypes.h"
-#include "llvm/IR/IRBuilder.h"
-#include "llvm/IR/Type.h"
-
-#include <algorithm>
-
-using namespace mlir;
-
-namespace {
-/// Utility function for type casting: this is making the type checker happy,
-/// while delaying the actual work involved to convert the type. Most of the
-/// time both side of the cast (producer and consumer) will be lowered to a
-/// dialect like LLVM and end up with the same LLVM representation, at which
-/// point this becomes a no-op and is eliminated.
-Value *typeCast(FuncBuilder &builder, Value *val, Type destTy) {
-  if (val->getType() == destTy)
-    return val;
-  return builder.create<toy::TypeCastOp>(val->getLoc(), val, destTy)
-      .getResult();
-}
-
-/// Create a type cast to turn a toy.array into a memref. The Toy Array will be
-/// lowered to a memref during buffer allocation, at which point the type cast
-/// becomes useless.
-Value *memRefTypeCast(FuncBuilder &builder, Value *val) {
-  if (val->getType().isa<MemRefType>())
-    return val;
-  auto toyArrayTy = val->getType().dyn_cast<toy::ToyArrayType>();
-  if (!toyArrayTy)
-    return val;
-  return typeCast(builder, val, toyArrayTy.toMemref());
-}
-
-/// Lower a toy.add to an affine loop nest.
-///
-/// This class inherit from `DialectOpConversion` and override `rewrite`,
-/// similarly to the PatternRewriter introduced in the previous chapter.
-/// It will be called by the DialectConversion framework (see `LateLowering`
-/// class below).
-class AddOpConversion : public DialectOpConversion {
-public:
-  explicit AddOpConversion(MLIRContext *context)
-      : DialectOpConversion(toy::AddOp::getOperationName(), 1, context) {}
-
-  /// Lower the `op` by generating IR using the `rewriter` builder. The builder
-  /// is setup with a new function, the `operands` array has been populated with
-  /// the rewritten operands for `op` in the new function.
-  /// The results created by the new IR with the builder are returned, and their
-  /// number must match the number of result of `op`.
-  SmallVector<Value *, 4> rewrite(Operation *op, ArrayRef<Value *> operands,
-                                  FuncBuilder &rewriter) const override {
-    auto add = op->cast<toy::AddOp>();
-    auto loc = add.getLoc();
-    // Create a `toy.alloc` operation to allocate the output buffer for this op.
-    Value *result = memRefTypeCast(
-        rewriter, rewriter.create<toy::AllocOp>(loc, add.getResult()->getType())
-                      .getResult());
-    Value *lhs = memRefTypeCast(rewriter, operands[0]);
-    Value *rhs = memRefTypeCast(rewriter, operands[1]);
-
-    using namespace edsc;
-    ScopedContext scope(rewriter, loc);
-    ValueHandle zero = intrinsics::constant_index(0);
-    MemRefView vRes(result), vLHS(lhs), vRHS(rhs);
-    IndexedValue iRes(result), iLHS(lhs), iRHS(rhs);
-    IndexHandle i, j, M(vRes.ub(0));
-    if (vRes.rank() == 1) {
-      LoopNestBuilder({&i}, {zero}, {M}, {1})({iRes(i) = iLHS(i) + iRHS(i)});
-    } else {
-      assert(vRes.rank() == 2 && "only rank 1 and 2 are supported right now");
-      IndexHandle N(vRes.ub(1));
-      LoopNestBuilder({&i, &j}, {zero, zero}, {M, N},
-                      {1, 1})({iRes(i, j) = iLHS(i, j) + iRHS(i, j)});
-    }
-
-    // Return the newly allocated buffer, with a type.cast to preserve the
-    // consumers.
-    return {typeCast(rewriter, result, add.getType())};
-  }
-};
-
-/// Lowers `toy.print` to a loop nest calling `printf` on every individual
-/// elements of the array.
-class PrintOpConversion : public DialectOpConversion {
-public:
-  explicit PrintOpConversion(MLIRContext *context)
-      : DialectOpConversion(toy::PrintOp::getOperationName(), 1, context) {}
-
-  SmallVector<Value *, 4> rewrite(Operation *op, ArrayRef<Value *> operands,
-                                  FuncBuilder &rewriter) const override {
-    // Get or create the declaration of the printf function in the module.
-    Function *printfFunc = getPrintf(*op->getFunction()->getModule());
-
-    auto print = op->cast<toy::PrintOp>();
-    auto loc = print.getLoc();
-    // We will operate on a MemRef abstraction, we use a type.cast to get one
-    // if our operand is still a Toy array.
-    Value *operand = memRefTypeCast(rewriter, operands[0]);
-    Type retTy = printfFunc->getType().getResult(0);
-
-    // Create our loop nest now
-    using namespace edsc;
-    using llvmCall = intrinsics::ValueBuilder<LLVM::CallOp>;
-    ScopedContext scope(rewriter, loc);
-    ValueHandle zero = intrinsics::constant_index(0);
-    ValueHandle fmtCst(getConstantCharBuffer(rewriter, loc, "%f "));
-    MemRefView vOp(operand);
-    IndexedValue iOp(operand);
-    IndexHandle i, j, M(vOp.ub(0));
-
-    ValueHandle fmtEol(getConstantCharBuffer(rewriter, loc, "\n"));
-    if (vOp.rank() == 1) {
-      // clang-format off
-      LoopBuilder(&i, zero, M, 1)({
-        llvmCall(retTy,
-                 rewriter.getFunctionAttr(printfFunc),
-                 {fmtCst, iOp(i)})
-      });
-      llvmCall(retTy, rewriter.getFunctionAttr(printfFunc), {fmtEol});
-      // clang-format on
-    } else {
-      IndexHandle N(vOp.ub(1));
-      // clang-format off
-      LoopBuilder(&i, zero, M, 1)({
-        LoopBuilder(&j, zero, N, 1)({
-          llvmCall(retTy,
-                   rewriter.getFunctionAttr(printfFunc),
-                   {fmtCst, iOp(i, j)})
-        }),
-        llvmCall(retTy, rewriter.getFunctionAttr(printfFunc), {fmtEol})
-      });
-      // clang-format on
-    }
-    return {};
-  }
-
-private:
-  // Turn a string into a toy.alloc (malloc/free abstraction) and a sequence
-  // of stores into the buffer, and return a MemRef into the buffer.
-  Value *getConstantCharBuffer(FuncBuilder &builder, Location loc,
-                               StringRef data) const {
-    auto retTy =
-        builder.getMemRefType(data.size() + 1, builder.getIntegerType(8));
-    Value *result = builder.create<toy::AllocOp>(loc, retTy).getResult();
-    using namespace edsc;
-    using intrinsics::constant_index;
-    using intrinsics::constant_int;
-    ScopedContext scope(builder, loc);
-    MemRefView vOp(result);
-    IndexedValue iOp(result);
-    for (uint64_t i = 0; i < data.size(); ++i) {
-      iOp(constant_index(i)) = constant_int(data[i], 8);
-    }
-    iOp(constant_index(data.size())) = constant_int(0, 8);
-    return result;
-  }
-
-  /// Return the prototype declaration for printf in the module, create it if
-  /// necessary.
-  Function *getPrintf(Module &module) const {
-    auto *printfFunc = module.getNamedFunction("printf");
-    if (printfFunc)
-      return printfFunc;
-
-    // Create a function declaration for printf, signature is `i32 (i8*, ...)`
-    Builder builder(&module);
-    MLIRContext *context = module.getContext();
-    LLVM::LLVMDialect *llvmDialect = static_cast<LLVM::LLVMDialect *>(
-        module.getContext()->getRegisteredDialect("llvm"));
-    auto &llvmModule = llvmDialect->getLLVMModule();
-    llvm::IRBuilder<> llvmBuilder(llvmModule.getContext());
-
-    auto llvmI32Ty = LLVM::LLVMType::get(context, llvmBuilder.getIntNTy(32));
-    auto llvmI8PtrTy =
-        LLVM::LLVMType::get(context, llvmBuilder.getIntNTy(8)->getPointerTo());
-    auto printfTy = builder.getFunctionType({llvmI8PtrTy}, {llvmI32Ty});
-    printfFunc = new Function(builder.getUnknownLoc(), "printf", printfTy);
-    // It should be variadic, but we don't support it fully just yet.
-    printfFunc->setAttr("std.varargs", builder.getBoolAttr(true));
-    module.getFunctions().push_back(printfFunc);
-    return printfFunc;
-  }
-};
-
-/// Lowers constant to a sequence of store in a buffer.
-class ConstantOpConversion : public DialectOpConversion {
-public:
-  explicit ConstantOpConversion(MLIRContext *context)
-      : DialectOpConversion(toy::ConstantOp::getOperationName(), 1, context) {}
-
-  SmallVector<Value *, 4> rewrite(Operation *op, ArrayRef<Value *> operands,
-                                  FuncBuilder &rewriter) const override {
-    toy::ConstantOp cstOp = op->cast<toy::ConstantOp>();
-    auto loc = cstOp.getLoc();
-    auto retTy = cstOp.getResult()->getType().cast<toy::ToyArrayType>();
-    auto shape = retTy.getShape();
-    Value *result = memRefTypeCast(
-        rewriter, rewriter.create<toy::AllocOp>(loc, retTy).getResult());
-
-    auto cstValue = cstOp.getValue();
-    auto f64Ty = rewriter.getF64Type();
-    using namespace edsc;
-    using intrinsics::constant_float;
-    using intrinsics::constant_index;
-    ScopedContext scope(rewriter, loc);
-    MemRefView vOp(result);
-    IndexedValue iOp(result);
-    for (uint64_t i = 0; i < shape[0]; ++i) {
-      if (shape.size() == 1) {
-        auto value = cstValue.getValue(ArrayRef<uint64_t>{i})
-                         .cast<FloatAttr>()
-                         .getValue();
-        iOp(constant_index(i)) = constant_float(value, f64Ty);
-        continue;
-      }
-      for (uint64_t j = 0; j < shape[1]; ++j) {
-        auto value = cstValue.getValue(ArrayRef<uint64_t>{i, j})
-                         .cast<FloatAttr>()
-                         .getValue();
-        iOp(constant_index(i), constant_index(j)) =
-            constant_float(value, f64Ty);
-      }
-    }
-    return {result};
-  }
-};
-
-/// Lower transpose operation to an affine loop nest.
-class TransposeOpConversion : public DialectOpConversion {
-public:
-  explicit TransposeOpConversion(MLIRContext *context)
-      : DialectOpConversion(toy::TransposeOp::getOperationName(), 1, context) {}
-
-  SmallVector<Value *, 4> rewrite(Operation *op, ArrayRef<Value *> operands,
-                                  FuncBuilder &rewriter) const override {
-    auto transpose = op->cast<toy::TransposeOp>();
-    auto loc = transpose.getLoc();
-    Value *result = memRefTypeCast(
-        rewriter,
-        rewriter.create<toy::AllocOp>(loc, transpose.getResult()->getType())
-            .getResult());
-    Value *operand = memRefTypeCast(rewriter, operands[0]);
-
-    using namespace edsc;
-    ScopedContext scope(rewriter, loc);
-    ValueHandle zero = intrinsics::constant_index(0);
-    MemRefView vRes(result), vOperand(operand);
-    IndexedValue iRes(result), iOperand(operand);
-    IndexHandle i, j, M(vRes.ub(0)), N(vRes.ub(1));
-    // clang-format off
-    LoopNestBuilder({&i, &j}, {zero, zero}, {M, N}, {1, 1})({
-      iRes(i, j) = iOperand(j, i)
-    });
-    // clang-format on
-
-    return {typeCast(rewriter, result, transpose.getType())};
-  }
-};
-
-// Lower toy.return to standard return operation.
-class ReturnOpConversion : public DialectOpConversion {
-public:
-  explicit ReturnOpConversion(MLIRContext *context)
-      : DialectOpConversion(toy::ReturnOp::getOperationName(), 1, context) {}
-
-  SmallVector<Value *, 4> rewrite(Operation *op, ArrayRef<Value *> operands,
-                                  FuncBuilder &rewriter) const override {
-    auto retOp = op->cast<toy::ReturnOp>();
-    using namespace edsc;
-    auto loc = retOp.getLoc();
-    // Argument is optional, handle both cases.
-    if (retOp.getNumOperands())
-      rewriter.create<ReturnOp>(loc, operands[0]);
-    else
-      rewriter.create<ReturnOp>(loc);
-    return {};
-  }
-};
-
-/// This is the main class registering our individual converter classes with
-/// the DialectConversion framework in MLIR.
-class LateLowering : public DialectConversion {
-protected:
-  /// Initialize the list of converters.
-  llvm::DenseSet<DialectOpConversion *>
-  initConverters(MLIRContext *context) override {
-    return ConversionListBuilder<AddOpConversion, PrintOpConversion,
-                                 ConstantOpConversion, TransposeOpConversion,
-                                 ReturnOpConversion>::build(&allocator,
-                                                            context);
-  }
-
-  /// Convert a Toy type, this gets called for block and region arguments, and
-  /// attributes.
-  Type convertType(Type t) override {
-    if (auto array = t.cast<toy::ToyArrayType>()) {
-      return array.toMemref();
-    }
-    return t;
-  }
-
-private:
-  llvm::BumpPtrAllocator allocator;
-};
-
-/// This is lowering to Linalg the parts that can be (matmul and add on arrays)
-/// and is targeting LLVM otherwise.
-struct LateLoweringPass : public ModulePass<LateLoweringPass> {
-
-  void runOnModule() override {
-    // Perform Toy specific lowering
-    if (failed(LateLowering().convert(&getModule()))) {
-      getModule().getContext()->emitError(
-          UnknownLoc::get(getModule().getContext()), "Error lowering Toy\n");
-      signalPassFailure();
-    }
-    // At this point the IR is almost using only standard and affine dialects.
-    // A few things remain before we emit LLVM IR. First to reuse as much of
-    // MLIR as possible we will try to lower everything to the standard and/or
-    // affine dialect: they already include conversion to the LLVM dialect.
-
-    // First patch calls type to return memref instead of ToyArray
-    for (auto &function : getModule()) {
-      function.walk([&](Operation *op) {
-        auto callOp = op->dyn_cast<CallOp>();
-        if (!callOp)
-          return;
-        if (!callOp.getNumResults())
-          return;
-        auto retToyTy =
-            callOp.getResult(0)->getType().dyn_cast<toy::ToyArrayType>();
-        if (!retToyTy)
-          return;
-        callOp.getResult(0)->setType(retToyTy.toMemref());
-      });
-    }
-
-    for (auto &function : getModule()) {
-      function.walk([&](Operation *op) {
-        // Turns toy.alloc into sequence of alloc/dealloc (later malloc/free).
-        if (auto allocOp = op->dyn_cast<toy::AllocOp>()) {
-          auto result = allocTensor(allocOp);
-          allocOp.replaceAllUsesWith(result);
-          allocOp.erase();
-          return;
-        }
-        // Eliminate all type.cast before lowering to LLVM.
-        if (auto typeCastOp = op->dyn_cast<toy::TypeCastOp>()) {
-          typeCastOp.replaceAllUsesWith(typeCastOp.getOperand());
-          typeCastOp.erase();
-          return;
-        }
-      });
-    }
-
-    // Lower Linalg to affine
-    for (auto &function : getModule())
-      linalg::lowerToLoops(&function);
-
-    getModule().dump();
-
-    // Finally convert to LLVM Dialect
-    linalg::convertLinalg3ToLLVM(getModule());
-  }
-
-  /// Allocate buffers (malloc/free) for Toy operations. This can't be done as
-  /// part of dialect conversion framework since we need to insert `dealloc`
-  /// operations just before the return, but the conversion framework is
-  /// operating in a brand new function: we don't have the return to hook the
-  /// dealloc operations.
-  Value *allocTensor(toy::AllocOp alloc) {
-    FuncBuilder builder(alloc);
-    auto retTy = alloc.getResult()->getType();
-
-    auto memRefTy = retTy.dyn_cast<MemRefType>();
-    if (!memRefTy)
-      memRefTy = retTy.cast<toy::ToyArrayType>().toMemref();
-    if (!memRefTy) {
-      alloc.emitOpError("is expected to allocate a Toy array or a MemRef");
-      llvm_unreachable("fatal error");
-    }
-    auto loc = alloc.getLoc();
-    Value *result = builder.create<AllocOp>(loc, memRefTy).getResult();
-
-    // Insert a `dealloc` operation right before the `return` operations, unless
-    // it is returned itself in which case the caller is responsible for it.
-    builder.getFunction()->walk([&](Operation *op) {
-      auto returnOp = op->dyn_cast<ReturnOp>();
-      if (!returnOp)
-        return;
-      if (returnOp.getNumOperands() && returnOp.getOperand(0) == alloc)
-        return;
-      builder.setInsertionPoint(returnOp);
-      builder.create<DeallocOp>(alloc.getLoc(), result);
-    });
-    return result;
-  }
-};
-} // end anonymous namespace
-
-namespace toy {
-Pass *createLateLoweringPass() { return new LateLoweringPass(); }
-
-std::unique_ptr<DialectConversion> makeToyLateLowering() {
-  return llvm::make_unique<LateLowering>();
-}
-
-} // namespace toy