[MLIR] Move mlir related classes to MLIR namespace (#23)

* Move dialect and types to mlir namespace

* PR fixes and some cleanup

* Merge fix

src/contrib/mlir/dialect/dialect.cpp

@@ -18,11 +18,8 @@
 #include "ops.hpp"
 #include "type.hpp"
 
-using namespace ngraph::runtime::ngmlir;
+using namespace mlir;
 
-/// Register a dialect and its types
-/// Usage:
-/// mlir::registerDialect<ngraph::runtime::ngmlir::Dialect>();
 NGDialect::NGDialect(mlir::MLIRContext* ctx)
     : mlir::Dialect("ng", ctx)
 {

src/contrib/mlir/dialect/dialect.hpp

@@ -23,24 +23,17 @@
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
 #include "ngraph/assertion.hpp"
-
-namespace ngraph
+namespace mlir
 {
-    namespace runtime
+    class NGDialect : public mlir::Dialect
     {
-        namespace ngmlir
+    public:
+        explicit NGDialect(mlir::MLIRContext* ctx);
+        mlir::Type parseType(llvm::StringRef tyData, mlir::Location loc) const override
         {
-            class NGDialect : public mlir::Dialect
-            {
-            public:
-                explicit NGDialect(mlir::MLIRContext* ctx);
-                mlir::Type parseType(llvm::StringRef tyData, mlir::Location loc) const override
-                {
-                    NGRAPH_ASSERT(0) << "Unsupported type parsing.";
-                    return mlir::Type();
-                }
-                void printType(mlir::Type type, llvm::raw_ostream& os) const override;
-            };
+            NGRAPH_ASSERT(0) << "Unsupported type parsing.";
+            return mlir::Type();
         }
-    }
+        void printType(mlir::Type type, llvm::raw_ostream& os) const override;
+    };
 }

src/contrib/mlir/dialect/ops.cpp

@@ -26,83 +26,69 @@ using llvm::raw_string_ostream;
 using llvm::SmallVector;
 using llvm::StringRef;
 using llvm::Twine;
-
-namespace ngraph
+using namespace mlir;
+// TODO:
+// - Move verifiers and other OP helpers (e.g. getSomeAttribute()) to separate files
+//
+// - Op helpers: Since it is not possible to add arbitrary code (and would complicate the .td file)
+// to Ops classes, we will add helper classes with static methods for each Op that needs it
+// Additional verification methods
+// Tensor type checks are already verified by the caller of these methods
+template <typename T>
+static mlir::LogicalResult verifyUnaryArithOp(T* op)
 {
-    namespace runtime
-    {
-        namespace ngmlir
-        {
-            // TODO:
-            // - Move verifiers and other OP helpers (e.g. getSomeAttribute()) to separate files
-            //
-            // - Op helpers: Since it is not possible to add arbitrary code (and would complicate the .td file)
-            // to Ops classes, we will add helper classes with static methods for each Op that needs it
-
-            // Additional verification methods
-            // Tensor type checks are already verified by the caller of these methods
-            template <typename T>
-            static mlir::LogicalResult verifyUnaryArithOp(T* op)
-            {
-                // TODO: Check matching element types
-                return mlir::success();
-            }
+    // TODO: Check matching element types
+    return mlir::success();
+}
 
-            // Additional verification methods
-            // Tensor type checks are already verified by the caller of these methods
-            template <typename T>
-            static mlir::LogicalResult verifyBinaryArithOp(T* op)
-            {
-                // TODO: Check matching element types
-                return mlir::success();
-            }
+// Additional verification methods
+// Tensor type checks are already verified by the caller of these methods
+template <typename T>
+static mlir::LogicalResult verifyBinaryArithOp(T* op)
+{
+    // TODO: Check matching element types
+    return mlir::success();
+}
 
-            template <typename T>
-            static mlir::LogicalResult verifyOp(T* op)
-            {
-                return op->emitOpError("Unsupported verifier for this operation");
-            }
+template <typename T>
+static mlir::LogicalResult verifyOp(T* op)
+{
+    return op->emitOpError("Unsupported verifier for this operation");
+}
 
-            // Per op specializations
-            template <>
-            mlir::LogicalResult verifyOp<NGMatMulBiasOp>(NGMatMulBiasOp* op)
-            {
-                // Verify that we have 2 operands
-                // Bias operand must be null for now (not implemented)
-                if (op->getNumOperands() != 2)
-                {
-                    std::stringstream ss;
-                    ss << "Unexpected MatmulBiasOp with " << op->getNumOperands()
-                       << " operands. 3 operands expected";
-                    return op->emitOpError(ss.str());
-                }
+// Per op specializations
+template <>
+mlir::LogicalResult verifyOp<NGMatMulBiasOp>(NGMatMulBiasOp* op)
+{
+    // Verify that we have 2 operands
+    // Bias operand must be null for now (not implemented)
+    if (op->getNumOperands() != 2)
+    {
+        std::stringstream ss;
+        ss << "Unexpected MatmulBiasOp with " << op->getNumOperands()
+           << " operands. 3 operands expected";
+        return op->emitOpError(ss.str());
+    }
 
-                // Verify that operand types are supported.
-                auto op0_tensor_ty = op->getOperand(0)->getType().cast<NGTensorType>();
-                auto op1_tensor_ty = op->getOperand(1)->getType().cast<NGTensorType>();
+    // Verify that operand types are supported.
+    auto op0_tensor_ty = op->getOperand(0)->getType().cast<NGTensorType>();
+    auto op1_tensor_ty = op->getOperand(1)->getType().cast<NGTensorType>();
 
-                // Verify that operand shapes are supported.
-                if (op0_tensor_ty.getRank() != 2 || op1_tensor_ty.getRank() != 2)
-                {
-                    return op->emitOpError(
-                        "Unsupported number of dimensions. Only 2D tensors are supported in "
-                        "MatmulBiasOp");
-                }
+    // Verify that operand shapes are supported.
+    if (op0_tensor_ty.getRank() != 2 || op1_tensor_ty.getRank() != 2)
+    {
+        return op->emitOpError(
+            "Unsupported number of dimensions. Only 2D tensors are supported in "
+            "MatmulBiasOp");
+    }
 
-                // TODO(dcab): Improve verification: matching types, proper shapes, etc.
+    // TODO(dcab): Improve verification: matching types, proper shapes, etc.
 
-                return mlir::success();
-            }
-        }
-    }
+    return mlir::success();
+}
 
-    using namespace mlir;
-    namespace runtime
-    {
-        namespace ngmlir
-        {
+namespace mlir
+{
 #define GET_OP_CLASSES
 #include "ops.cpp.inc"
-        }
-    }
 }

src/contrib/mlir/dialect/ops.hpp

@@ -22,19 +22,8 @@
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/STLExtras.h"
 
-namespace ngraph
+namespace mlir
 {
-    namespace runtime
-    {
-        namespace ngmlir
-        {
-            // TODO: We shouldn't have this here, but we need to expose mlir types for the .inc file to use
-            // we cannot forward declare the mlir types since they rely on the Ops we are defining (see. Op<NGAddOp, ...>)
-            //
-            // Other ways to avoid namespace pollution ?
-            using namespace mlir;
 #define GET_OP_CLASSES
 #include "ops.h.inc"
-        }
-    }
 }

src/contrib/mlir/dialect/ops.td

@@ -40,7 +40,7 @@ include "mlir/IR/OpBase.td"
 // This defines records equivalent to NGraph types. It doesn't generate code.
 // This is used as a type in the DAG input/outputs.
 // Constraints (CPred) are used to type-check args/results of that type during op verification
-def NG_TensorType : Type<CPred<"{0}.isa<ngraph::runtime::ngmlir::NGTensorType>()">,
+def NG_TensorType : Type<CPred<"{0}.isa<mlir::NGTensorType>()">,
                      "NGraph Tensor Type">;
 
 // A generic un-typed MemRef. Used for Fake instructions inserted during dialect lowering

src/contrib/mlir/dialect/type.cpp

@@ -31,55 +31,51 @@ using llvm::SmallVector;
 using llvm::StringRef;
 using llvm::Twine;
 
-namespace ngraph
-{
-    using namespace runtime::ngmlir;
+using namespace mlir;
 
-    unsigned NGIntegerType::getWidth() const
+unsigned NGIntegerType::getWidth() const
+{
+    switch (getKind())
     {
-        switch (getKind())
-        {
-        case NG_I8_TYPE_ID:
-        case NG_U8_TYPE_ID: return 8;
-        case NG_I16_TYPE_ID:
-        case NG_U16_TYPE_ID: return 16;
-        case NG_I32_TYPE_ID:
-        case NG_U32_TYPE_ID: return 32;
-        case NG_I64_TYPE_ID:
-        case NG_U64_TYPE_ID: return 64;
-        default: NGRAPH_FAIL() << "Invalid type ID";
-        }
-        return 0;
+    case NG_I8_TYPE_ID:
+    case NG_U8_TYPE_ID: return 8;
+    case NG_I16_TYPE_ID:
+    case NG_U16_TYPE_ID: return 16;
+    case NG_I32_TYPE_ID:
+    case NG_U32_TYPE_ID: return 32;
+    case NG_I64_TYPE_ID:
+    case NG_U64_TYPE_ID: return 64;
+    default: NGRAPH_FAIL() << "Invalid type ID";
     }
+    return 0;
+}
 
-    bool NGIntegerType::isSigned() const
+bool NGIntegerType::isSigned() const
+{
+    switch (getKind())
     {
-        switch (getKind())
-        {
-        case NG_I8_TYPE_ID:
-        case NG_I16_TYPE_ID:
-        case NG_I32_TYPE_ID:
-        case NG_I64_TYPE_ID: return true;
-        case NG_U8_TYPE_ID:
-        case NG_U16_TYPE_ID:
-        case NG_U32_TYPE_ID:
-        case NG_U64_TYPE_ID: return false;
-        default: NGRAPH_FAIL() << "Invalid type ID";
-        }
-        return false;
+    case NG_I8_TYPE_ID:
+    case NG_I16_TYPE_ID:
+    case NG_I32_TYPE_ID:
+    case NG_I64_TYPE_ID: return true;
+    case NG_U8_TYPE_ID:
+    case NG_U16_TYPE_ID:
+    case NG_U32_TYPE_ID:
+    case NG_U64_TYPE_ID: return false;
+    default: NGRAPH_FAIL() << "Invalid type ID";
     }
+    return false;
+}
 
-    /// Creates TensorType objects. They all point to the same storage if
-    /// element type and shape are the same.
-    NGTensorType NGTensorType::get(mlir::MLIRContext* context, EltType eltType, Shape shape)
-    {
-        return Base::get(context, NGTypeKind::NG_TENSOR_TYPE_ID, eltType, shape);
-    }
+/// Creates TensorType objects. They all point to the same storage if
+/// element type and shape are the same.
+NGTensorType NGTensorType::get(MLIRContext* context, EltType eltType, Shape shape)
+{
+    return Base::get(context, NGTypeKind::NG_TENSOR_TYPE_ID, eltType, shape);
+}
 
-    mlir::MemRefType NGTensorType::toMemref()
-    {
-        auto memRefType =
-            mlir::MemRefType::get(getShape(), getElementType(), {/* no map used */}, 0);
-        return memRefType;
-    }
+MemRefType NGTensorType::toMemref()
+{
+    auto memRefType = MemRefType::get(getShape(), getElementType(), {/* no map used */}, 0);
+    return memRefType;
 }

src/contrib/mlir/lowerer.cpp

@@ -131,7 +131,7 @@ namespace
 
         // we find out output values by looking at returned values
         // any return should return all outputs of the subgraph
-        f->walk<ngmlir::NGReturnOp>([this, &outputCount](ngmlir::NGReturnOp ret) {
+        f->walk<NGReturnOp>([this, &outputCount](NGReturnOp ret) {
             for (unsigned i = 0; i < ret.getNumOperands(); i++)
             {
                 this->m_outputValueMap.insert(std::pair<Value*, unsigned>(ret.getOperand(i), i));
@@ -151,8 +151,8 @@ namespace
         // however, due to how DialectConversion framework works, new func is only
         // materialized after conversion is done (rewriter->getFunction, or even rewriter->getInsertionBlock()->getFunction()
         // will give you the original func). This makes it very convoluted to insert instructions at entry block.
-        auto op = rewriter->create<ngmlir::NGFakeInputOp>(rewriter->getUnknownLoc(),
-                                                          IndexType::get(getModule().getContext()));
+        auto op = rewriter->create<NGFakeInputOp>(rewriter->getUnknownLoc(),
+                                                  IndexType::get(getModule().getContext()));
         // will be fixed later to read passed arg instead.
         m_memMgrDefs.push_back(op.getResult());
         return op.getResult();
@@ -170,7 +170,7 @@ namespace
             if (it != outputMap.end())
             {
                 unsigned argId = (*it).second;
-                auto fakeOp = rewriter.create<ngmlir::NGFakeInputOp>(
+                auto fakeOp = rewriter.create<NGFakeInputOp>(
                     op->getLoc(),
                     m_dialectLowerer.convertType(
                         origResult->getType()) /* convert to lowered type */
@@ -183,7 +183,7 @@ namespace
             }
             else
             {
-                auto tensorType = origResult->getType().cast<ngmlir::NGTensorType>();
+                auto tensorType = origResult->getType().cast<NGTensorType>();
                 auto callBackFunc = getCallDecl("__mlir_allocate",
                                                 {rewriter.getIndexType(), rewriter.getIndexType()},
                                                 {tensorType.toMemref()},
@@ -237,8 +237,7 @@ namespace
         for (auto value : m_loweredOutputValues)
         {
             auto op = value->getDefiningOp();
-            NGRAPH_ASSERT(op->isa<ngmlir::NGFakeInputOp>())
-                << "output value not defined by fake output?";
+            NGRAPH_ASSERT(op->isa<NGFakeInputOp>()) << "output value not defined by fake output?";
             value->replaceAllUsesWith(entryBlock->getArgument(oldFuncType.getNumInputs() + i));
             op->erase();
             i++;
@@ -269,23 +268,23 @@ namespace
     // NGDialect converters
     Type DialectLowerer::convertType(Type t)
     {
-        if (auto tensor = t.dyn_cast<ngmlir::NGTensorType>())
+        if (auto tensor = t.dyn_cast<NGTensorType>())
         {
             return tensor.toMemref();
         }
         // element type
-        if (auto type = t.dyn_cast<ngmlir::NGFloatType>())
+        if (auto type = t.dyn_cast<NGFloatType>())
         {
             // Float
             // float types are already std type
             return type;
         }
-        if (auto type = t.dyn_cast<ngmlir::NGIntegerType>())
+        if (auto type = t.dyn_cast<NGIntegerType>())
         {
             // map it to std type
             return type.toStdType();
         }
-        if (auto type = t.dyn_cast<ngmlir::NGBoolType>())
+        if (auto type = t.dyn_cast<NGBoolType>())
         {
             return type.toStdType();
         }
@@ -298,7 +297,7 @@ namespace
                                                       ArrayRef<Value*> operands,
                                                       FuncBuilder& rewriter) const
     {
-        auto add = op->cast<ngmlir::NGAddOp>();
+        auto add = op->cast<NGAddOp>();
         auto loc = add.getLoc();
         Value *origResult, *newResult;
 
@@ -335,7 +334,7 @@ namespace
                                                              ArrayRef<Value*> operands,
                                                              FuncBuilder& rewriter) const
     {
-        auto matmul = op->cast<ngmlir::NGMatMulBiasOp>();
+        auto matmul = op->cast<NGMatMulBiasOp>();
         auto loc = matmul.getLoc();
 
         NGRAPH_ASSERT(operands.size() == 2) << "Bias is not supported yet in MatmulBias operation";
@@ -406,16 +405,10 @@ namespace
     }
 }
 
-namespace ngraph
+namespace mlir
 {
-    namespace runtime
+    Pass* createDialectLoweringPass(ngraph::runtime::ngmlir::MLIRCompiler* compiler)
     {
-        namespace ngmlir
-        {
-            Pass* createDialectLoweringPass(MLIRCompiler* compiler)
-            {
-                return new DialectLoweringPass(*compiler);
-            }
-        }
+        return new DialectLoweringPass(*compiler);
     }
 }

src/contrib/mlir/lowerer.hpp

@@ -16,9 +16,9 @@
 
 #pragma once
 
+#include "contrib/mlir/compiler.hpp"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
-
 namespace ngraph
 {
     namespace runtime
@@ -26,8 +26,10 @@ namespace ngraph
         namespace ngmlir
         {
             class MLIRCompiler;
-
-            mlir::Pass* createDialectLoweringPass(MLIRCompiler* compiler);
         }
     }
 }
+namespace mlir
+{
+    mlir::Pass* createDialectLoweringPass(ngraph::runtime::ngmlir::MLIRCompiler* compiler);
+}

src/contrib/mlir/op_lowerers.inc

@@ -22,7 +22,7 @@ class OP##Conversion : public mlir::DialectOpConversion \
 {\
 public:\
     explicit OP##Conversion(mlir::MLIRContext *context, DialectLoweringPass& pass)\
-    : mlir::DialectOpConversion(ngraph::runtime::ngmlir::OP::getOperationName(), 1, context),\
+    : mlir::DialectOpConversion(mlir::OP::getOperationName(), 1, context),\
       m_pass(pass)\
     {} \
     SmallVector<Value *, 4> rewrite(Operation *op, ArrayRef<Value *> operands, FuncBuilder &rewriter) const override; \