Merge pull request #3263 from NervanaSystems/nmostafa/gather

[MLIR] Enable Gather Op

src/contrib/mlir/compiler.hpp

@@ -98,25 +98,21 @@ namespace ngraph
 
                 void build_ng_dialect();
 
-                template <typename OP>
-                static mlir::Value* create_op(MLIRCompiler& compiler, const ngraph::Node* ng_node)
+                template <typename Op>
+                static mlir::Operation* create_op(MLIRCompiler& compiler,
+                                                  const ngraph::Node* ng_node)
                 {
                     throw std::runtime_error("Unimplemented op '" + ng_node->description() +
                                              "' in MLIR Compiler");
                 }
 
-                template <typename UnaryOp>
-                mlir::Value* create_unary_op(const ngraph::Node* ng_node);
-
-                template <typename BinOp>
-                mlir::Value* create_binary_op(const ngraph::Node* ng_node);
-
-                // TODO(amprocte): Can we have a create_variadic_op that is able to handle the
-                // attributes?
-                mlir::Value* create_concat(const ngraph::Node* ng_node);
+                // Generic op lowerer to ng dialect.
+                // Simply maps ngraph tensors to values and generate an OP. No op-specific logic.
+                template <typename Op>
+                mlir::Operation* create_generic_op(const ngraph::Node* ng_node);
 
                 template <typename RedOp>
-                mlir::Value* create_index_reduction(const ngraph::Node* ng_node);
+                mlir::Operation* create_index_reduction(const ngraph::Node* ng_node);
 
                 void create_return();
 
@@ -150,7 +146,7 @@ namespace ngraph
                 using TensorToInfo = std::pair<descriptor::Tensor*, TensorInfo>;
                 using TensorToInfoMap = std::unordered_map<descriptor::Tensor*, TensorInfo>;
                 using MLIRCompOpFunction =
-                    std::function<mlir::Value*(MLIRCompiler& compiler, const ngraph::Node*)>;
+                    std::function<mlir::Operation*(MLIRCompiler& compiler, const ngraph::Node*)>;
                 using MLIRCompOpMap = std::unordered_map<std::type_index, MLIRCompOpFunction>;
 
                 // Maps tensor to the value it represents in the IR

src/contrib/mlir/dialect/ops.cpp

@@ -168,6 +168,39 @@ static mlir::LogicalResult verifyCmpOp(T* op)
     return mlir::success();
 }
 
+template <>
+mlir::LogicalResult verifyOp(NGGatherOp* op)
+{
+    Type ty = op->params()->getType();
+    NGTensorType inputType = ty.cast<NGTensorType>();
+
+    ty = op->indices()->getType();
+    NGTensorType indicesType = ty.cast<NGTensorType>();
+
+    // ensure axis < params rank
+    if (op->axis().getSExtValue() >= inputType.getRank())
+        return op->emitOpError("Gather axis is larger than input rank");
+
+    ty = indicesType.getElementType();
+
+    // ensure indices are I32 or I64
+    if (!ty.isa<NGIntegerType>())
+        return op->emitOpError("Indices tensor is not of Integer type");
+
+    NGIntegerType indicesEltType = ty.cast<NGIntegerType>();
+    if (!indicesEltType.isInt32() && !indicesEltType.isInt64())
+        return op->emitOpError("Indices tensor is not of I32 or I64 type");
+
+    mlir::Type r0 = op->res()->getType();
+    NGTensorType resType = r0.cast<NGTensorType>();
+
+    // ensure result is compatible with input
+    if (!resType.getRank() == inputType.getRank() + indicesType.getRank() - 1)
+        return op->emitOpError("Incompatible result shape and/or type");
+
+    return mlir::success();
+}
+
 namespace mlir
 {
 #define GET_OP_CLASSES

src/contrib/mlir/dialect/ops.td

@@ -186,8 +186,8 @@ def NGDotOp : NG_Binary_Op<"dot">
 // class, but I'm not sure how to add concatenation_axis into the args if we
 // do that.
 def NGConcatOp :
-      NG_OneResult_Op<"concat", [NoSideEffect]>,
-      Arguments<(ins Variadic<NG_TensorType>:$args, I64Attr:$concatenation_axis)>
+    NG_OneResult_Op<"concat", [NoSideEffect]>,
+    Arguments<(ins Variadic<NG_TensorType>:$args, I64Attr:$concatenation_axis)>
 {
   let parser = [{ NGRAPH_CHECK(false, "No parser support"); return mlir::failure(); }];
 
@@ -200,7 +200,7 @@ class NG_Axis_Reduction_Op<string mnemonic, list<OpTrait> traits = []> :
 {
   let summary = "Base class for reduction operations that perform a reduction "
                 "across the axes of a  single tensor.";
-  let description = "Axes are represented as an array of I64 attributes.";
+  let description = [{Axes are represented as an array of I64 attributes.}];
 
   let parser = [{ NGRAPH_CHECK(false, "No parser support"); return mlir::failure(); }];
 
@@ -257,6 +257,24 @@ def NGAnyRedOp : NG_Axis_Reduction_Op<"any.red">
   let verifier = [{ return verifyLogicalReductionOp(this); }];
 }
 
+// Gather
+def NGGatherOp : 
+    NG_OneResult_Op<"gather", [NoSideEffect]>,
+    Arguments<(ins NG_TensorType:$params, NG_TensorType:$indices, I64Attr:$axis)>
+{
+  let summary = "Gather slices from params along the specified axis according to indices";
+  let description = [{
+    Gather slices from axis of params according to indices
+    params The tensor from which slices are gathered
+    indices Index tensor. Data type must be `element::i32` or `element::i64`
+    axis Axis in params to gather
+  }];
+
+  let parser = [{ NGRAPH_CHECK(false, "No parser support"); return mlir::failure(); }];
+
+  let verifier = [{ return verifyOp(this); }];
+}
+
 // Terminator Ops
 def NGReturnOp : NG_Terminator_Op<"return">;
 

src/contrib/mlir/dialect/type.hpp

@@ -199,6 +199,7 @@ namespace mlir
         }
 
         Shape getShape() const { return m_shape; }
+        int64_t getRank() const { return m_shape.size(); }
         EltType getElementType() const { return m_eltType; }
     private:
         NGTensorTypeStorage(EltType eltType, Shape shape)

src/contrib/mlir/lowerer.cpp

@@ -646,6 +646,123 @@ namespace
         return matchSuccess();
     }
 
+    REWRITER(NGGatherOp)
+    {
+        auto gatherOp = cast<NGGatherOp>(op);
+        auto loc = gatherOp.getLoc();
+        ScopedContext scope(rewriter, loc);
+
+        // Get operands
+        Value* result = m_pass.buildOutputDefs(op, rewriter)[0];
+        NGRAPH_CHECK(result, "Unexpected null result in GatherOp");
+        auto resultTy = result->getType().cast<MemRefType>();
+
+        Value* params = operands[0];
+        Value* indices = operands[1];
+        auto axis = gatherOp.axis().getSExtValue();
+
+        // Create view to write into result.
+        MemRefView vRes(result), vParams(params), vIndices(indices);
+        // Indexed Values
+        IndexedValue iRes(result), iParams(params), iIndices(indices);
+
+        // Construct outer loop for params dims. Exclude the axis dim.
+        SmallVector<ValueHandle, 4> paramsLbs, paramsUbs;
+        SmallVector<IndexHandle, 4> paramsIVs;
+        SmallVector<int64_t, 4> paramsSteps;
+        SmallVector<ValueHandle*, 4> paramsIVPtrs;
+        for (auto i = 0; i < vParams.rank(); i++)
+        {
+            // skip gather axis
+            if (i == axis)
+                continue;
+            paramsLbs.push_back(IndexHandle(vParams.lb(i)));
+            paramsUbs.push_back(IndexHandle(vParams.ub(i)));
+            paramsSteps.push_back(vParams.step(i));
+        }
+        NGRAPH_CHECK(paramsLbs.size() == vParams.rank() - 1 &&
+                         paramsUbs.size() == paramsLbs.size() &&
+                         paramsSteps.size() == paramsLbs.size(),
+                     "Incorrect loop nest bounds size for gather params");
+
+        paramsIVs = IndexHandle::makeIndexHandles(vParams.rank() - 1);
+        paramsIVPtrs = IndexHandle::makeIndexHandlePointers(paramsIVs);
+
+        auto indicesLbs = vIndices.getLbs();
+        auto indicesUbs = vIndices.getUbs();
+        auto indicesSteps = vIndices.getSteps();
+
+        auto indicesIVs = IndexHandle::makeIndexHandles(vIndices.rank());
+        auto indicesIVPtrs = IndexHandle::makeIndexHandlePointers(indicesIVs);
+
+        SmallVector<IndexHandle, 8> paramsIndices, resIndices;
+
+        // Make sure we are going to create loops
+        NGRAPH_CHECK(vParams.rank() > 0, "Invalid size for indices steps");
+
+        // Let params rank : N
+        // Let indices rank : M
+        // Let axis be A
+        // Generate
+        // params loops
+        // for P_0: 0 -> params.dim[0]
+        //   for P_1: 0 -> params.dim[1]
+        //     for P_2: 0 -> params.dim[2]
+        // ...
+        //       for P_(A-1):0 -> params.dim[A-1]
+        //         for P_(A+1):0 -> params.dim[A+1]
+        // ...
+        //           for P_(N-1):0 -> params.dim[N-1]
+        //             indices loops
+        //             for I_0:0 -> indices.dim[0]
+        // ...
+        //               for I_(M-1):0 -> indices.dim[M-1]
+        //                 res[P_0, P_1, .. P_(A-1), I_0, .., I_(M-1), P_(A+1), ... P_(N-1)] =
+        //                   params[P_0, P_1, .. P_(A-1), indices[I_0, .., I_(M-1)], P_(A+1), ... P_(N-1)];
+
+        LoopNestBuilder(paramsIVPtrs, paramsLbs, paramsUbs, paramsSteps)([&] {
+            LoopNestBuilder(indicesIVPtrs, indicesLbs, indicesUbs, indicesSteps)([&] {
+                // Load axis value from indices array and cast it to Index Type
+                ValueHandle axisIdx = ValueHandle::create<IndexCastOp>(
+                    (ValueHandle)iIndices(indicesIVs), rewriter.getIndexType());
+                // construct indices for param
+                // [P_0, P_1, .. P_axis-1, Indices[I0, I1, .. I_k-1], P_axis+1, P_axis+2, .. P_n-1]
+                for (auto i = 0, j = 0; i < vParams.rank(); i++)
+                {
+                    if (i == axis)
+                    {
+                        paramsIndices.push_back(IndexHandle(axisIdx));
+                    }
+                    else
+                    {
+                        paramsIndices.push_back(paramsIVs[j++]);
+                    }
+                }
+
+                // construct indices for result
+                // [P_0, P_1, .. P_axis-1, I0, I1, .. I_k-1, P_axis+1, P_axis+2, .. P_n-1]
+                for (auto i = 0, j = 0; i < vParams.rank() + vIndices.rank() - 1;)
+                {
+                    if (i == axis && indicesIVs.size() > 0)
+                    {
+                        resIndices.append(indicesIVs.begin(), indicesIVs.end());
+                        i += indicesIVs.size();
+                    }
+                    else
+                    {
+                        resIndices.push_back(paramsIVs[j++]);
+                        i++;
+                    }
+                }
+                // Store into result
+                iRes(resIndices) = iParams(paramsIndices);
+            });
+        });
+
+        rewriter.replaceOp(op, {result});
+        return matchSuccess();
+    }
+
     REWRITER(NGReturnOp)
     {
         rewriter.replaceOpWithNewOp<ReturnOp>(op);
@@ -653,7 +770,7 @@ namespace
     }
 
 #undef REWRITER
-
+    /// End of pattern matchers
     template <typename OP>
     void lower_binary_elementwise(Operation* op,
                                   ArrayRef<Value*> operands,

src/contrib/mlir/op_lowerers.inc

@@ -29,6 +29,7 @@ MLIR_OP(NGArgMinRedOp)
 MLIR_OP(NGConcatOp)
 MLIR_OP(NGDivOp)
 MLIR_OP(NGDotOp)
+MLIR_OP(NGGatherOp)
 MLIR_OP(NGGreaterOp)
 MLIR_OP(NGLessOp)
 MLIR_OP(NGMulOp)

src/contrib/mlir/ops_supported.inc

@@ -9,6 +9,7 @@ MLIR_OP(ArgMax)
 MLIR_OP(Divide)
 MLIR_OP(Dot)
 MLIR_OP(Concat)
+MLIR_OP(Gather)
 MLIR_OP(Greater)
 MLIR_OP(Less)
 MLIR_OP(Maximum)

src/contrib/mlir/pass/mlir_subgraph_extraction.cpp

@@ -25,6 +25,7 @@
 #include "ngraph/op/divide.hpp"
 #include "ngraph/op/dot.hpp"
 #include "ngraph/op/experimental/compiled_kernel.hpp"
+#include "ngraph/op/gather.hpp"
 #include "ngraph/op/get_output_element.hpp"
 #include "ngraph/op/greater.hpp"
 #include "ngraph/op/less.hpp"