[Intel] Remove legacy DPAS conversion (#3529)

This PR removes `ConvertLayoutOpConversion`, which is the legacy way of
converting DPAS layouts. The pass now relies on Linear Layout to perform
such conversions.

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

third_party/intel/lib/TritonIntelGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -31,162 +31,6 @@ namespace {
 // because LinearLayout seems to have some performance issues.
 constexpr bool useLegacyMMAConversion = false;
 
-struct ConvertLayoutOpConversion
-    : public ConvertTritonGPUOpToLLVMPattern<triton::gpu::ConvertLayoutOp> {
-public:
-  ConvertLayoutOpConversion(const LLVMTypeConverter &typeConverter,
-                            const triton::intel::TargetInfo &targetInfo,
-                            PatternBenefit benefit = 1)
-      : ConvertTritonGPUOpToLLVMPattern(typeConverter, benefit) {}
-
-  LogicalResult
-  matchAndRewrite(triton::gpu::ConvertLayoutOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    RankedTensorType srcTy = op.getSrc().getType();
-    RankedTensorType dstTy = op.getType();
-    Attribute srcLayout = srcTy.getEncoding();
-    Attribute dstLayout = dstTy.getEncoding();
-    if (isa<DpasEncodingAttr>(srcLayout) &&
-        isa<DotOperandEncodingAttr>(dstLayout)) {
-      return lowerDpasToDotOperand(op, adaptor, rewriter);
-    }
-    return failure();
-  }
-
-private:
-  using ValueTable = std::map<std::array<unsigned, 3>, Value>;
-
-  ValueTable getValuesFromDpasLayoutStruct(Location loc,
-                                           ConversionPatternRewriter &rewriter,
-                                           Value vals,
-                                           RankedTensorType srcType) const {
-    SmallVector<Value> elems = unpackLLElements(loc, vals, rewriter);
-    auto dpasLayout = dyn_cast<DpasEncodingAttr>(srcType.getEncoding());
-
-    size_t totalElems = elems.size();
-    auto numElemsPerOperand =
-        product<unsigned>(dpasLayout.getDPASInstShapeC()) /
-        product<unsigned>(dpasLayout.getThreadsPerWarp());
-    Type elemTy =
-        this->getTypeConverter()->convertType(srcType.getElementType());
-    VectorType dotOpTy = vec_ty(elemTy, numElemsPerOperand);
-    SmallVector<int64_t> repetitions =
-        dpasLayout.getDPASRepetitions(srcType.getShape(), 2 /*operand C*/);
-    ArrayRef<unsigned> repCluster = dpasLayout.getRepCluster();
-    size_t rank = repCluster.size();
-    size_t outerDim = rank - 2;
-    size_t innerDim = rank - 1;
-
-    auto tb = TritonLLVMOpBuilder(loc, rewriter);
-    int offset = 0;
-    ValueTable result;
-    for (unsigned b = 0; b < repetitions[0]; ++b) {
-      for (int i = 0; i < repetitions[1]; ++i) {
-        for (int j = 0; j < repetitions[2]; ++j) {
-          for (int repOuter = 0; repOuter < repCluster[outerDim]; ++repOuter) {
-            for (int repInner = 0; repInner < repCluster[innerDim];
-                 ++repInner) {
-              Value matVal = rewriter.create<LLVM::UndefOp>(loc, dotOpTy);
-              for (int k = 0; k < numElemsPerOperand; ++k) {
-                matVal = tb.insert_element(dotOpTy, matVal, elems[offset++],
-                                           tb.i32_val(k));
-              }
-              result[{b, i * repCluster[outerDim] + repOuter,
-                      j * repCluster[innerDim] + repInner}] = matVal;
-            }
-          }
-        }
-      }
-    }
-
-    return result;
-  }
-
-  Value composeValuesToDotOperandLayoutStruct(
-      Location loc, ConversionPatternRewriter &rewriter, const ValueTable &vals,
-      RankedTensorType dstType) const {
-    auto tb = TritonLLVMOpBuilder(loc, rewriter);
-    auto dotLayout = dyn_cast<DotOperandEncodingAttr>(dstType.getEncoding());
-    auto dpasLayout = dyn_cast<DpasEncodingAttr>(dotLayout.getParent());
-
-    auto opIdx = static_cast<DpasEncodingAttr::OpIdx>(dotLayout.getOpIdx());
-    SmallVector<int64_t> repetitions =
-        dpasLayout.getDPASRepetitions(dstType.getShape(), opIdx);
-    ArrayRef<unsigned> repCluster = dpasLayout.getRepCluster();
-    size_t rank = repCluster.size();
-    unsigned repBatch = repetitions[0];
-    unsigned repOuter = 0u;
-    unsigned repInner = 0u;
-    unsigned repClusterOuter = 0u;
-
-    switch (opIdx) {
-    case DpasEncodingAttr::OpIdx::OperandA: {
-      // operand A
-      repOuter = repetitions[1];
-      repInner = repetitions[2];
-      repClusterOuter = repCluster[rank - 2];
-    } break;
-    case DpasEncodingAttr::OpIdx::OperandB: {
-      // operand B
-      repOuter = repetitions[2];
-      repInner = repetitions[1];
-      repClusterOuter = repCluster[rank - 1];
-    } break;
-    case DpasEncodingAttr::OpIdx::OperandC: {
-      llvm_unreachable("unexpected OpIdx::OperandC");
-    } break;
-    }
-
-    // TODO: Operands B requires extra steps to combine [8, 16] to [16, 16].
-    SmallVector<Value> elems;
-    for (unsigned b = 0; b < repBatch; ++b) {
-      for (int m = 0; m < repOuter; ++m) {
-        for (int k = 0; k < repInner; ++k) {
-          for (int repOuterIdx = 0; repOuterIdx < repClusterOuter;
-               ++repOuterIdx) {
-            unsigned offsetM = m * repClusterOuter + repOuterIdx;
-            unsigned offsetN = k;
-            Value matVal = vals.at({b, offsetM, offsetN});
-            auto vecType = cast<VectorType>(matVal.getType());
-            Type valTy = vecType.getElementType();
-            for (int i = 0; i < vecType.getNumElements(); ++i) {
-              Value val = tb.extract_element(valTy, matVal, tb.i32_val(i));
-              elems.push_back(val);
-            }
-          }
-        }
-      }
-    }
-
-    Type elemTy = getTypeConverter()->convertType(dstType.getElementType());
-    Type structTy = LLVM::LLVMStructType::getLiteral(
-        getContext(), SmallVector<Type>(elems.size(), elemTy));
-    return packLLElements(loc, this->getTypeConverter(), elems, rewriter,
-                          structTy);
-  }
-
-  // dpas -> dot_operand
-  LogicalResult
-  lowerDpasToDotOperand(triton::gpu::ConvertLayoutOp op, OpAdaptor adaptor,
-                        ConversionPatternRewriter &rewriter) const {
-    Location loc = op.getLoc();
-    RankedTensorType srcTy = op.getSrc().getType();
-    RankedTensorType dstTy = op.getType();
-
-    if (!intel::isDpasToDotShortcut(srcTy, dstTy))
-      return failure();
-
-    // reorder the elements to match the dot_operand layout.
-    ValueTable values =
-        getValuesFromDpasLayoutStruct(loc, rewriter, adaptor.getSrc(), srcTy);
-    Value view =
-        composeValuesToDotOperandLayoutStruct(loc, rewriter, values, dstTy);
-
-    rewriter.replaceOp(op, view);
-    return success();
-  }
-};
-
 struct ConvertLayoutOpUsingLinearLayoutsConversion
     : public ConvertOpToLLVMPattern<ConvertLayoutOp> {
   const TargetInfoBase &targetInfo;
@@ -1006,8 +850,6 @@ void mlir::triton::intel::populateConvertLayoutOpToLLVMPatterns(
   // and be the only one left.
   patterns.add<gpu::ConvertLayoutOpUsingLinearLayoutsConversion>(
       typeConverter, targetInfo, benefit.getBenefit() + 2);
-  patterns.add<gpu::ConvertLayoutOpConversion>(typeConverter, targetInfo,
-                                               benefit.getBenefit() + 1);
   mlir::triton::populateConvertLayoutOpToLLVMPatterns(typeConverter, targetInfo,
                                                       patterns, benefit);
 }