[HINT] [Shared_Memory] Support hint the use of shared memory for aipu (#29)

* [CONVERSION] Support the use of shared memory for AIPU.

* When using shared memory, the annotation "@hint: shared_memory" is required.
* MemorySpace value 8 of shared memory is only for AIPU backend.

* [CONVERSION] Fix shared memory for aipu.
 test_02_fused_softmax.py requires support later.

* [Conversion] Add out-of-bounds offset control for shared memory

 At present in the use of Shared memory, need to comment out mlir/lib/the Dialect/Affine/Analysis/LoopAnalysis.cpp
 in the following two lines, and recompile LLVM:

 //if (! memRefType.getLayout().isIdentity())
 //  return memoryOp.emitError("NYI: non-trivial layout map"), false;

* [Conversion] Fix the Cases of Softmax Kernel and One-Dimensional Non-Mask.

* Supports two-dimensional matrix addition using shared memory ;
 Fixed the bug that caused errors in the operation of 2D matrix addition.

* [Conversion] Support multi-dimensional vector add and  add the "shared" attribute.
 * The extension supports multi-dimensional vector add.
 * Add the "shared" attribute to the memref.subview carrying offset of shared memory
   to solve the problem of non-trivial layout.

* [Conversion] Downward propagation of the 'shared_memory' hint in memref.alloc

* [Conversion] Add the no_bufferize attribute to memref ops with shared memory address space

---------

Co-authored-by: CHAIYD-2025 <chaiyd_z@163.com>
Co-authored-by: duyb <2697533527@qq.com>

Author: 3 people

include/triton-shared/Conversion/CMakeLists.txt

@@ -7,3 +7,4 @@ add_subdirectory(TritonToUnstructured)
 add_subdirectory(StructuredToMemref)
 add_subdirectory(UnstructuredToMemref)
 add_subdirectory(MemrefCopyToDMA_FlagTree)
+add_subdirectory(NoBufferize_FlagTree)

include/triton-shared/Conversion/NoBufferize_FlagTree/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls --name NoBufferizeFlagTree)
+add_public_tablegen_target(NoBufferizeFlagTreeConversionPassIncGen)

include/triton-shared/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.h

@@ -0,0 +1,24 @@
+#ifndef TRITON_CONVERSION_NOBUFFERIZEFLAGTREE_NOBUFFERIZEFLAGTREE_H
+#define TRITON_CONVERSION_NOBUFFERIZEFLAGTREE_NOBUFFERIZEFLAGTREE_H
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+namespace mlir {
+class TypeConverter;
+namespace triton {
+
+#define GEN_PASS_DECL
+#include "triton-shared/Conversion/NoBufferize_FlagTree/Passes.h.inc"
+
+void populateNoBufferizeFlagTreeConversionPatterns(
+    RewritePatternSet &patterns, TypeConverter &typeConverter);
+
+std::unique_ptr<OperationPass<ModuleOp>> createNoBufferizeFlagTreePass();
+
+} // namespace triton
+} // namespace mlir
+
+#endif // TRITON_CONVERSION_NOBUFFERIZEFLAGTREE_NOBUFFERIZEFLAGTREE_H

include/triton-shared/Conversion/NoBufferize_FlagTree/Passes.h

@@ -0,0 +1,15 @@
+#ifndef TRITON_NO_BUFFERIZE_FLAGTREE_CONVERSION_PASSES_H
+#define TRITON_NO_BUFFERIZE_FLAGTREE_CONVERSION_PASSES_H
+
+#include "triton-shared/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.h"
+
+namespace mlir {
+namespace triton {
+
+#define GEN_PASS_REGISTRATION
+#include "triton-shared/Conversion/NoBufferize_FlagTree/Passes.h.inc"
+
+} // namespace triton
+} // namespace mlir
+
+#endif

include/triton-shared/Conversion/NoBufferize_FlagTree/Passes.td

@@ -0,0 +1,10 @@
+#ifndef NO_BUFFERIZE_FLAGTREE_CONVERSION_PASSES
+#define NO_BUFFERIZE_FLAGTREE_CONVERSION_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def NoBufferizeFlagTree : Pass<"no-bufferize-flagtree", "mlir::ModuleOp"> {
+  let summary = "Add no_bufferize to memref ops in shared memory space";
+}
+
+#endif

lib/AnalysisStructured/PtrAnalysis.cpp

@@ -1160,6 +1160,10 @@ LogicalResult PtrAnalysis::rewriteLoadOp(triton::LoadOp op,
     loadOp->setAttr("flagtree_hints", strAttr);
   }
 
+  if (op->getAttr("flagtree_hints")) {
+    loadOp->setAttr("flagtree_hints", op->getAttr("flagtree_hints"));
+  }
+
   LLVM_DEBUG({
     llvm::dbgs() << "creating tts::load:\n";
     loadOp->dump();

lib/Conversion/CMakeLists.txt

@@ -7,3 +7,4 @@ add_subdirectory(StructuredToMemref)
 add_subdirectory(TritonPtrToMemref)
 add_subdirectory(UnstructuredToMemref)
 add_subdirectory(MemrefCopyToDMA_FlagTree)
+add_subdirectory(NoBufferize_FlagTree)

lib/Conversion/NoBufferize_FlagTree/CMakeLists.txt

@@ -0,0 +1,24 @@
+add_triton_library(NoBufferizeFlagTree
+  NoBufferizeFlagTree.cpp
+  NoBufferizeFlagTreePass.cpp
+
+  DEPENDS
+  NoBufferizeFlagTreeConversionPassIncGen
+  Registrar
+  Common
+
+  LINK_LIBS PUBLIC
+  MLIRSCFTransforms
+  MLIRArithDialect
+  MLIRDialectUtils
+  MLIRIR
+  MLIRMathDialect
+  MLIRPass
+  MLIRTensorDialect
+  MLIRTransforms
+  MLIRSupport
+  TritonIR
+  TritonTransforms
+  TritonTilingExtIR
+  TritonStructuredIR
+)

lib/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.cpp

@@ -0,0 +1,65 @@
+#include "flagtree/Common/UnifiedHardware.h"
+
+#include "triton/Dialect/Triton/IR/Types.h"
+
+#include "triton-shared/Analysis/OpFoldResultUtils.h"
+#include "triton-shared/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.h"
+#include "triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <iostream>
+
+#define DEBUG_TYPE "no-bufferize-flagtree"
+
+using namespace mlir;
+#define GEN_PASS_CLASSES
+#include "triton-shared/Conversion/NoBufferize_FlagTree/Passes.h.inc"
+#include "triton-shared/Conversion/TritonArithToLinalg/Passes.h.inc"
+
+namespace {
+struct NoBufferizeConverter : public RewritePattern {
+  NoBufferizeConverter(MLIRContext *ctx)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    auto memrefNoBufferize = [](Type t) -> bool {
+      Attribute msAttr;
+      if (auto memTy = dyn_cast<MemRefType>(t))
+        msAttr = memTy.getMemorySpace();
+      else if (auto unranked = dyn_cast<UnrankedMemRefType>(t))
+        msAttr = unranked.getMemorySpace();
+      else
+        return false;
+      if (auto intAttr = dyn_cast_or_null<IntegerAttr>(msAttr))
+        return intAttr.getValue().getZExtValue() == 8;
+      return false;
+    };
+
+    bool noBuf = llvm::any_of(op->getResultTypes(), memrefNoBufferize);
+    if (!noBuf)
+      return failure();
+
+    op->setAttr("no_bufferize", BoolAttr::get(op->getContext(), true));
+    return success();
+  }
+};
+} // namespace
+
+void mlir::triton::populateNoBufferizeFlagTreeConversionPatterns(
+    RewritePatternSet &patterns, TypeConverter &typeConverter) {
+  patterns.add<NoBufferizeConverter>(patterns.getContext());
+}

lib/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTreePass.cpp

@@ -0,0 +1,64 @@
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+#include "triton-shared/Conversion/NoBufferize_FlagTree/NoBufferizeFlagTree.h"
+#include "triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.h"
+#include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/Transforms/Patterns.h"
+#include "mlir/Pass/PassManager.h"
+#include "triton/Dialect/Triton/IR/Types.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include <iostream>
+#include <optional>
+
+#define DEBUG_TYPE "no-bufferize-flagtree"
+using namespace mlir;
+using namespace triton;
+
+namespace mlir {
+namespace triton {
+#define GEN_PASS_DEF_NOBUFFERIZEFLAGTREE
+#include "triton-shared/Conversion/NoBufferize_FlagTree/Passes.h.inc"
+} // namespace triton
+} // namespace mlir
+
+namespace {
+
+class NoBufferizeFlagTreePass
+    : public triton::impl::NoBufferizeFlagTreeBase<NoBufferizeFlagTreePass> {
+  using NoBufferizeFlagTreeBase<
+      NoBufferizeFlagTreePass>::NoBufferizeFlagTreeBase;
+
+public:
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    RewritePatternSet localPatterns(&getContext());
+    {
+      TypeConverter typeConverter;
+      triton::populateNoBufferizeFlagTreeConversionPatterns(localPatterns,
+                                                            typeConverter);
+    }
+    (void)applyPatternsGreedily(module, std::move(localPatterns));
+  }
+};
+} // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>>
+triton::createNoBufferizeFlagTreePass() {
+  return std::make_unique<NoBufferizeFlagTreePass>();
+}