DMA, Support hint for triton dsl

lib/AnalysisStructured/PtrAnalysis.cpp

@@ -1155,6 +1155,10 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op,
   }
 
   auto loadOp = builder.create<tts::LoadOp>(loc, ptr, dims, scalarOther);
+  auto strAttr = op->getAttrOfType<mlir::StringAttr>("flagtree_hints");
+  if (strAttr && !strAttr.getValue().empty()) {
+    loadOp->setAttr("flagtree_hints", strAttr);
+  }
 
   LLVM_DEBUG({
     llvm::dbgs() << "creating tts::load:\n";

lib/Conversion/MemrefCopyToDMA_FlagTree/MemrefCopyToDMAFlagTree.cpp

@@ -84,11 +84,12 @@ struct CopyConverter : public OpConversionPattern<memref::CopyOp> {
 
   LogicalResult rewriteCopyToDma(memref::CopyOp op, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
-    
+
     auto hardwareManager = mlir::flagtree::createUnifiedHardwareManager();
-    auto dmaTag = hardwareManager -> getDMATag();
-    if (!dmaTag) return failure();
-
+    auto dmaTag = hardwareManager->getDMATag();
+    if (!dmaTag)
+      return failure();
+
     Location loc = op.getLoc();
     Value src = adaptor.getSource();
     Value dst = adaptor.getTarget();
@@ -122,19 +123,16 @@ struct CopyConverter : public OpConversionPattern<memref::CopyOp> {
     //       * A synchronization buffer of type memref<1xi32, *>.
     //         The memory space `*` denotes a hardware-reserved region for DMA
     //         completion signaling, determined by the unified hardware layer.
+    int64_t rank = mlir::cast<MemRefType>(src.getType()).getRank();
+    srcIndices.assign(rank, zero);
+    dstIndices.assign(rank, zero);
 
     if (auto srcSubview = dyn_cast_or_null<memref::SubViewOp>(srcDef)) {
       auto dstSubview = dyn_cast<memref::SubViewOp>(dst.getDefiningOp());
 
-      srcIndices = getValueList(rewriter, loc, srcSubview.getMixedOffsets());
-      dstIndices = getValueList(rewriter, loc, dstSubview.getMixedOffsets());
       auto sizes = getValueList(rewriter, loc, srcSubview.getMixedSizes());
       numElements = getTotalElementCount(rewriter, loc, sizes);
     } else if (auto castOp = dyn_cast<memref::ReinterpretCastOp>(srcDef)) {
-      int64_t rank = mlir::cast<MemRefType>(src.getType()).getRank();
-
-      srcIndices.assign(rank, zero);
-      dstIndices.assign(rank, zero);
       auto sizes = getValueList(rewriter, loc, castOp.getMixedSizes());
       numElements = getTotalElementCount(rewriter, loc, sizes);
     } else {
@@ -163,9 +161,11 @@ struct CopyConverter : public OpConversionPattern<memref::CopyOp> {
   matchAndRewrite(memref::CopyOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
 
-    // TODO: default to convert all memref.copy to memref.dma. Future will
-    // support hint for triton dsl, relative configuration will in unified
-    // hardware.
+    auto strAttr = op->getAttrOfType<mlir::StringAttr>("flagtree_hints");
+    if (!strAttr || strAttr.getValue() != "dma") {
+      return success();
+    }
+
     bool isStrideOne = false;
     Value src = adaptor.getSource();
     Operation *srcDef = src.getDefiningOp();

lib/Conversion/StructuredToMemref/StructuredToMemref.cpp

@@ -599,7 +599,11 @@ struct LoadConverter : public OpConversionPattern<tts::LoadOp> {
         llvm_unreachable("unexpected wraparound type");
       }
     } else {
-      rewriter.create<memref::CopyOp>(loc, ptr, alloc);
+      auto copyOp = rewriter.create<memref::CopyOp>(loc, ptr, alloc);
+      auto strAttr = op->getAttrOfType<mlir::StringAttr>("flagtree_hints");
+      if (strAttr && !strAttr.getValue().empty()) {
+        copyOp->setAttr("flagtree_hints", strAttr);
+      }
     }
 
     Value tensor = rewriter.create<bufferization::ToTensorOp>(
@@ -686,7 +690,12 @@ struct LoadConverter : public OpConversionPattern<tts::LoadOp> {
           getSubview(tensorType.getRank(), mixedDims, ptr, loc, rewriter);
       memref::SubViewOp dstSubview =
           getSubview(tensorType.getRank(), mixedDims, alloc, loc, rewriter);
-      rewriter.create<memref::CopyOp>(loc, srcSubview, dstSubview);
+      auto copyOp =
+          rewriter.create<memref::CopyOp>(loc, srcSubview, dstSubview);
+      auto strAttr = op->getAttrOfType<mlir::StringAttr>("flagtree_hints");
+      if (strAttr && !strAttr.getValue().empty()) {
+        copyOp->setAttr("flagtree_hints", strAttr);
+      }
     }
 
     Value tensor = rewriter.create<bufferization::ToTensorOp>(

lib/Conversion/TritonToLinalgExperimental/TritonToLinalgExperimentalPass.cpp

@@ -4,6 +4,7 @@
 // Licensed under the MIT license.
 //
 //===----------------------------------------------------------------------===//
+#include "flagtree/Common/UnifiedHardware.h"
 
 #include "mlir/Dialect/Ptr/IR/PtrDialect.h"
 #include "mlir-ext/Dialect/MathExt/IR/MathExt.h"
@@ -75,9 +76,11 @@ class TritonToLinalgExperimentalPass
     // leave as a TODO for now.
     pm.addPass(createStructuredToMemrefPass());
 
-    // TODO: determine whether go this pass in some way, maybe in unified
-    // haredware.
-    pm.addPass(createMemrefCopyToDMAFlagTreePass());
+    // Pass selection is controlled by unified hardware configuration.
+    auto hardwareManager = mlir::flagtree::createUnifiedHardwareManager();
+    auto dmaTag = hardwareManager -> getDMATag();
+    if (dmaTag)
+      pm.addPass(createMemrefCopyToDMAFlagTreePass());
 
     pm.addPass(createUnstructuredToMemrefPass());
     pm.addPass(createTritonPtrToMemrefPass());