[FRONTEND] Added ExpandDimsOp primitive (#36)

lib/Analysis/AxisInfo.cpp

@@ -167,8 +167,8 @@ ChangeResult AxisInfoAnalysis::visitOperation(
     }
     curr = AxisInfo(contiguity, divisibility, constancy);
   }
-  // Reshape
-  if (llvm::isa<triton::ViewOp>(op)) {
+  // expandDims
+  if (auto expandDims = llvm::dyn_cast<triton::ExpandDimsOp>(op)) {
     Type _retTy = *op->result_type_begin();
     Type _opTy = *op->operand_type_begin();
     TensorType retTy = _retTy.cast<TensorType>();
@@ -176,28 +176,12 @@ ChangeResult AxisInfoAnalysis::visitOperation(
     ArrayRef<int64_t> retShape = retTy.getShape();
     ArrayRef<int64_t> opShape = opTy.getShape();
     AxisInfo opInfo = operands[0]->getValue();
-    AxisInfo::DimVectorT contiguity;
-    AxisInfo::DimVectorT divisibility;
-    AxisInfo::DimVectorT constancy;
-    bool is_skewed = false;
-    size_t current = 0;
-    for (size_t d = 0; d < retTy.getRank(); d++) {
-      if (retShape[d] == 1) {
-        contiguity.push_back(1);
-        divisibility.push_back(1);
-        constancy.push_back(1);
-      } else if (!is_skewed && retShape[d] == opShape[current]) {
-        contiguity.push_back(opInfo.getContiguity()[current]);
-        divisibility.push_back(opInfo.getDivisibility()[current]);
-        constancy.push_back(opInfo.getConstancy()[current]);
-        current++;
-      } else {
-        is_skewed = true;
-        contiguity.push_back(1);
-        divisibility.push_back(1);
-        constancy.push_back(1);
-      }
-    }
+    AxisInfo::DimVectorT contiguity = opInfo.getContiguity();
+    AxisInfo::DimVectorT divisibility = opInfo.getDivisibility();
+    AxisInfo::DimVectorT constancy = opInfo.getConstancy();
+    contiguity.insert(contiguity.begin() + expandDims.axis(), 1);
+    divisibility.insert(divisibility.begin() + expandDims.axis(), 1);
+    constancy.insert(constancy.begin() + expandDims.axis(), 1);
     curr = AxisInfo(contiguity, divisibility, constancy);
   }
   // Broadcast

lib/Conversion/TritonToTritonGPU/TritonToTritonGPU.cpp

@@ -5,6 +5,7 @@
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h"
+#include <numeric>
 
 using namespace mlir;
 using namespace mlir::triton;
@@ -142,6 +143,46 @@ struct TritonMakeRangePattern
   }
 };
 
+struct TritonExpandDimsPattern
+    : public OpConversionPattern<triton::ExpandDimsOp> {
+  using OpConversionPattern<triton::ExpandDimsOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(triton::ExpandDimsOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Type retType = op.getType());
+    RankedTensorType argType = adaptor.src().getType().cast<RankedTensorType>();
+    Attribute _argEncoding = argType.getEncoding();
+    if (!_argEncoding)
+      return failure();
+    auto argEncoding =
+        _argEncoding.cast<triton::gpu::TritonGPUBlockedEncodingAttr>();
+    // return shape
+    auto retShape = argType.getShape().vec();
+    retShape.insert(retShape.begin() + op.axis(), 1);
+    // return encoding
+    auto retSizePerThread = argEncoding.getSizePerThread().vec();
+    retSizePerThread.insert(retSizePerThread.begin() + op.axis(), 1);
+    auto retThreadsPerWarp = argEncoding.getThreadsPerWarp().vec();
+    retThreadsPerWarp.insert(retThreadsPerWarp.begin() + op.axis(), 1);
+    auto retWarpsPerCTA = argEncoding.getWarpsPerCTA().vec();
+    retWarpsPerCTA.insert(retWarpsPerCTA.begin() + op.axis(), 1);
+    SmallVector<unsigned, 4> retOrder(retShape.size());
+    std::iota(retOrder.begin(), retOrder.end(), 0);
+    triton::gpu::TritonGPUBlockedEncodingAttr retEncoding =
+        triton::gpu::TritonGPUBlockedEncodingAttr::get(
+            getContext(), retSizePerThread, retThreadsPerWarp, retWarpsPerCTA,
+            retOrder);
+    // return type
+    RankedTensorType retType =
+        RankedTensorType::get(retShape, argType.getElementType(), retEncoding);
+    // construct new op
+    rewriter.replaceOpWithNewOp<triton::ExpandDimsOp>(
+        op, retType, adaptor.src(), adaptor.axis());
+    return success();
+  }
+};
+
 struct TritonDotPattern : public OpConversionPattern<triton::DotOp> {
   using OpConversionPattern<triton::DotOp>::OpConversionPattern;
 
@@ -260,8 +301,8 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
       TritonGenericPattern<triton::ViewOp>,
       TritonGenericPattern<triton::SplatOp>, TritonBroadcastPattern,
       TritonGenericPattern<triton::GEPOp>, TritonReducePattern,
-      TritonMakeRangePattern, TritonDotPattern, TritonLoadPattern,
-      TritonStorePattern>(typeConverter, context);
+      TritonExpandDimsPattern, TritonMakeRangePattern, TritonDotPattern,
+      TritonLoadPattern, TritonStorePattern>(typeConverter, context);
 }
 
 //