[WIP][BACKEND][PPU] Add PPU backend

.gitignore

@@ -12,7 +12,7 @@ python/dist/
 python/_deps/
 python/triton*.egg-info/
 python/triton_kernels/triton*.egg-info/
-python/flagtree.egg-info/
+python/flagtree*.egg-info/
 
 python/triton/_C/*.pyd
 python/triton/_C/*.so

CMakeLists.txt

@@ -71,6 +71,8 @@ elseif(FLAGTREE_BACKEND STREQUAL "metax")
   list(REMOVE_ITEM LLVM_TABLEGEN_FLAGS -D__TLE__)
 elseif(FLAGTREE_BACKEND STREQUAL "sunrise")
   find_package(Python3 3.10 REQUIRED COMPONENTS Development.Module Interpreter)
+elseif(FLAGTREE_BACKEND STREQUAL "ppu")
+  add_definitions(-D__PPU__)
 endif()
 set(FLAGTREE_PLUGIN "$ENV{FLAGTREE_PLUGIN}")
 if(FLAGTREE_PLUGIN)
@@ -299,7 +301,7 @@ endfunction()
 # Disable warnings that show up in external code (gtest;pybind11)
 if(NOT MSVC)
   if(FLAGTREE_BACKEND)
-    if(FLAGTREE_BACKEND MATCHES "^(enflame|hcu|rpu|thrive|metax|xpu|tileir)$")
+    if(FLAGTREE_BACKEND MATCHES "^(enflame|hcu|rpu|thrive|metax|xpu|tileir|ppu)$")
       # Suppress visibility warnings in gluon_ir.cc (GCC 13+ -Wattributes on pybind11 hidden types)
       # Also suppress -Wcomment for generated TritonGPUAttrDefs.h.inc (ASCII diagrams in TableGen output)
       set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror -Wno-attributes -Wno-comment -Wno-error=odr")
@@ -336,7 +338,7 @@ if(BUILD_MCTLE)
 endif()
 
 # link_directories(${LLVM_LIBRARY_DIR})
-if (FLAGTREE_BACKEND MATCHES "^(xpu|cambricon|aipu|tsingmicro|enflame|rpu|thrive|tileir)$")
+if (FLAGTREE_BACKEND MATCHES "^(xpu|cambricon|aipu|tsingmicro|enflame|rpu|thrive|tileir|ppu)$")
   include_directories(${PROJECT_SOURCE_DIR}/include)
   include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files
   if(FLAGTREE_BACKEND STREQUAL "xpu")
@@ -764,7 +766,7 @@ find_package(Threads REQUIRED)
 
 add_subdirectory(third_party/f2reduce)
 
-if(NOT FLAGTREE_BACKEND OR FLAGTREE_BACKEND MATCHES "^(aipu|tsingmicro|enflame|rpu|thrive|metax|sunrise|tileir)$")
+if(NOT FLAGTREE_BACKEND OR FLAGTREE_BACKEND MATCHES "^(aipu|tsingmicro|enflame|rpu|thrive|metax|sunrise|tileir|ppu)$")
   add_subdirectory(bin)
   if(FLAGTREE_TLE)
     flagtree_add_tle_generated_header_dependencies()
@@ -778,6 +780,15 @@ elseif(FLAGTREE_BACKEND STREQUAL "iluvatar")
   endif()
 endif()
 
+# When PPU backend is built, the upstream TritonGPU transforms reference
+# PPU-specific dialect symbols (e.g. mlir::triton::ppu_gpu::AsyncAIUCopyGlobalToLocalOp)
+# that live in TritonPPUGPUIR (defined under third_party/ppu/lib/...). Wire that
+# dependency here, after both targets have been declared by their respective
+# add_subdirectory() calls above, so unittests/triton.so/triton-llvm-opt all link cleanly.
+if(FLAGTREE_BACKEND STREQUAL "ppu" AND TARGET TritonGPUTransforms AND TARGET TritonPPUGPUIR)
+  target_link_libraries(TritonGPUTransforms PUBLIC TritonPPUGPUIR)
+endif()
+
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)
   # This target runs all the unit tests.

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -516,6 +516,7 @@ using ::mlir::triton::gpu::AMDWmmaEncodingAttr;
 using ::mlir::triton::gpu::BlockedEncodingAttr;
 using ::mlir::triton::gpu::DotOperandEncodingAttr;
 using ::mlir::triton::gpu::NvidiaMmaEncodingAttr;
+using ::mlir::triton::gpu::PPUMmaEncodingAttr;
 using ::mlir::triton::gpu::SliceEncodingAttr;
 
 Value dot(RewriterBase &rewriter, Location loc, ArrayRef<Value> offsets,

include/triton/Dialect/Triton/IR/TritonOps.td

@@ -1104,7 +1104,8 @@ def TT_MakeTensorDescOp : TT_Op<"make_tensor_descriptor", [
     TT_Ptr:$base,
     Variadic<I32>:$shape,
     Variadic<I64>:$strides,
-    DefaultValuedAttr<TT_PaddingOptionAttr, "::mlir::triton::PaddingOption::PAD_ZERO">:$padding
+    DefaultValuedAttr<TT_PaddingOptionAttr, "::mlir::triton::PaddingOption::PAD_ZERO">:$padding,
+    DefaultValuedAttr<BoolAttr, "false">:$isAIU
   );
 
   let results = (outs TT_TensorDescType:$result);
@@ -1113,7 +1114,7 @@ def TT_MakeTensorDescOp : TT_Op<"make_tensor_descriptor", [
 
   let builders = [
     OpBuilder<(ins "Value":$base, "ValueRange":$shape, "ValueRange":$strides, "ArrayRef<int32_t>":$blockShape, "bool":$isSignedInteger,
-    "triton::PaddingOption":$padding)>
+    "triton::PaddingOption":$padding, CArg<"bool", "false">:$isAIU)>
   ];
 
   let extraClassDeclaration = [{
@@ -1123,6 +1124,41 @@ def TT_MakeTensorDescOp : TT_Op<"make_tensor_descriptor", [
   }];
 }
 
+def TT_AIULoadOp : TT_Op<"aiu_load", [
+  SameVariadicOperandSize,
+  MemoryEffects<[MemRead<GlobalMemory>]>]> {
+    let summary = "AIU Load";
+    let description = [{
+      This operation will be lowered to PPU AIU load operation on targets supporting it.
+    }];
+    let arguments = (
+      ins
+      AnyTypeOf<[TT_PtrLike, TT_TensorPtr]>:$src_ptr,
+      Variadic<I32>:$indices,
+      Variadic<I64>:$shape,
+      DefaultValuedAttr<DenseI32ArrayAttr, "::llvm::ArrayRef<int32_t>{1, 0}">:$order,
+      DefaultValuedAttr<TT_CacheModifierAttr, "::mlir::triton::CacheModifier::NONE">:$cache,
+      DefaultValuedAttr<TT_EvictionPolicyAttr, "::mlir::triton::EvictionPolicy::NORMAL">:$evict
+    );
+
+    let builders = [
+      // A tensor pointer with boundary check and padding
+      OpBuilder<(ins "Type":$type, "Value":$src_ptr, "triton::CacheModifier":$cache,
+                 "triton::EvictionPolicy":$evict)>,
+    ];
+
+    let results = (outs TT_Tensor:$result);
+
+    let assemblyFormat = [{
+      $src_ptr `[` $indices `]` `[` $shape `]` `[` $order `]`
+      oilist(
+        `cacheModifier` `=` $cache |
+        `evictionPolicy` `=` $evict
+      )
+      attr-dict `:` qualified(type($src_ptr)) `->` type($result)
+    }];
+}
+
 // The following ops, including `call`, `func`, and `return` are copied and modified from
 // https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
 // We could revert it back once MLIR has a better inliner interface.

include/triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h

@@ -142,6 +142,12 @@ chooseDsReadTrLayout(Attribute enc, ArrayRef<int64_t> shape,
                      int32_t elemBitWidth, unsigned instBitWidth,
                      unsigned numLanesInShuffleGroup);
 
+// The primary goal of this function is to efficiently load 2D tiles of a
+// tensor from shared memory using the `ldmatrix` instruction.
+LinearLayout choosePPULdMatrixLayout(Attribute enc, ArrayRef<int64_t> shape,
+                                     bool needTrans, int32_t elemBitWidth,
+                                     bool Opb8bLdmatrix);
+
 // Create LinearLayout for scale in scaled mfma.
 LinearLayout chooseScaledMfmaScaleLayout(MLIRContext *ctx, int dotOperandIdx,
                                          ArrayRef<int64_t> dotOperandShape,

include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td

@@ -193,6 +193,12 @@ When vec=2, elements are swizzled in pairs of 2.  In other words, the element at
               typeWidthInBit, needTrans);
         }
 
+        // ---- begin PPUMma ----
+        if (auto PPUMmaEnc = mlir::dyn_cast<PPUMmaEncodingAttr>(dotOpEnc.getParent())) {
+          if (PPUMmaEnc.isPPU0010() || PPUMmaEnc.isPPU0015()) {
+            return get(context, dotOpEnc.getOpIdx(), dotOpEnc.getKWidth(), shape, order, CTALayout, typeWidthInBit, needTrans);
+          }
+        }
 
         auto mmaEnc = mlir::dyn_cast<NvidiaMmaEncodingAttr>(dotOpEnc.getParent());
 
@@ -582,6 +588,49 @@ Swizzling examples (matrix is filled with numbers 0, 1, 2, .. columns*rows-1):
   let hasCustomAssemblyFormat = 1;
 }
 
+def PPUAIUSharedEncodingAttr :
+  TritonGPU_Attr<"PPUAIUSharedEncoding", "ppu_aiu_shared_encoding",
+                 [SharedEncodingTrait, LayoutEncodingTrait,
+                  DeclareLayoutEncodingMethods]> {
+  let mnemonic = "ppu_aiu_shared";
+
+  let description = [{
+    An encoding for tensors whose elements may be simultaneously accessed by
+    different threads in the programs, via shared memory.
+  }];
+
+  let parameters = (
+    ins
+    DefaultValuedParameter<"unsigned", "1">:$versionMajor,
+    ArrayRefParameter<"unsigned">:$AIUStrategy,
+    ArrayRefParameter<"unsigned">:$order,
+    "CTAEncodingAttr":$CTALayout,
+    DefaultValuedParameter<"unsigned", "0">:$kOffset
+  );
+
+  let builders = [
+    AttrBuilder<(ins "ArrayRef<unsigned>":$AIUStrategy,
+                     "ArrayRef<unsigned>":$order,
+                     "CTAEncodingAttr":$CTALayout), [{
+      //default kOffset=0;
+      return $_get(context, 0, AIUStrategy, order, CTALayout, 0);
+    }]>,
+    AttrBuilder<(ins "unsigned":$versionMajor,
+                     "ArrayRef<unsigned>":$AIUStrategy,
+                     "ArrayRef<unsigned>":$order,
+                     "CTAEncodingAttr":$CTALayout), [{
+      //default kOffset=0;
+      return $_get(context, versionMajor, AIUStrategy, order, CTALayout, 0);
+    }]>,
+  ];
+
+  let extraClassDeclaration = extraBaseClassDeclaration # [{
+    bool isPPU0010() const;
+    bool isPPU0015() const;
+  }];
+  let hasCustomAssemblyFormat = 1;
+}
+
 
 //===----------------------------------------------------------------------===//
 // Distributed Layout Encoding
@@ -1366,6 +1415,55 @@ For example, the matrix L corresponding to blockTileSize=[32,16] is:
   let hasCustomAssemblyFormat = 1;
 }
 
+def PPUMmaEncodingAttr : DistributedEncoding<"PPUMmaEncoding", "ppu_mma_encoding", [MmaEncodingTrait]> {
+  let mnemonic = "ppu_mma";
+
+  let description = [{
+An encoding for tensors that have been produced by tensor cores.
+
+It is characterized by two parameters:
+- A 'versionMajor' which specifies the generation the tensor cores
+  whose output is being partitioned:
+  - 1 for first-gen tensor cores (PPU0010), and
+  - 2 for second-gen tensor cores (PPU0015).
+- A 'versionMinor' which indicates the specific layout of a tensor core generation.
+- A `blockTileSize` to indicate how data should be partitioned between warps.
+}];
+
+  let parameters = (
+    ins
+    "unsigned":$versionMajor,
+    "unsigned":$versionMinor,
+    ArrayRefParameter<"unsigned">:$warpsPerCTA,
+    "CTAEncodingAttr":$CTALayout,
+    ArrayRefParameter<"unsigned">:$instrShape,
+    DefaultValuedParameter<"unsigned", "2">:$vecSize
+  );
+
+  let builders = [
+    AttrBuilder<(ins "unsigned":$versionMajor,
+                     "unsigned":$versionMinor,
+                     "ArrayRef<unsigned>":$warpsPerCTA,
+                     "CTAEncodingAttr":$CTALayout,
+                     "ArrayRef<unsigned>":$instrShape), [{
+                return $_get(context, versionMajor, versionMinor,
+                  warpsPerCTA, CTALayout, instrShape, 2);
+              }]>
+  ];
+
+  let extraClassDeclaration = extraDistributedDeclaration # [{
+    bool isPPU0010() const;
+    bool isPPU0015() const;
+
+    SmallVector<int64_t> getRepForOperand(ArrayRef<int64_t> shape,
+                                          int bitwidth, int kWidth,
+                                          int opIdx) const;
+    SmallVector<unsigned> getRepOrderForOperand(int opIdx) const;
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+}
+
 def SliceEncodingAttr : DistributedEncoding<"SliceEncoding", "slice_encoding"> {
   let mnemonic = "slice";
 
@@ -1455,20 +1553,41 @@ vecIdx (index of the element in the quad; this is always along the k-dim)
     ins
     "unsigned":$opIdx,
     "Attribute":$parent,
-    DefaultValuedParameter<"unsigned", "0">:$kWidth
+    DefaultValuedParameter<"unsigned", "0">:$kWidth,
+    DefaultValuedParameter<"bool", "false">:$isChained
   );
 
   let builders = [
     AttrBuilder<(ins "unsigned":$opIdx,
                      "Attribute":$parent,
-                     "Type":$eltTy), [{
+                     "unsigned":$kWidth), [{
+      return $_get(context, opIdx, parent, kWidth, false);
+    }]>,
+
+    AttrBuilder<(ins "unsigned":$opIdx,
+                     "Attribute":$parent,
+                     "Type":$eltTy,
+                     "bool":$isChained), [{
+      PPUMmaEncodingAttr PPUParentAttr = mlir::dyn_cast<PPUMmaEncodingAttr>(parent);
+      if (PPUParentAttr && (PPUParentAttr.isPPU0010() || PPUParentAttr.isPPU0015())) {
+        unsigned bitwidth = eltTy.getIntOrFloatBitWidth();
+        unsigned kWidth = std::max(32 / bitwidth, 1u);
+        return $_get(context, opIdx, parent, kWidth, isChained);
+      }
+
       NvidiaMmaEncodingAttr parentAttr = mlir::dyn_cast<NvidiaMmaEncodingAttr>(parent);
       if (!parentAttr || (!parentAttr.isAmpere() && !parentAttr.isHopper()))
-        return $_get(context, opIdx, parent, 0);
+        return $_get(context, opIdx, parent, 0, isChained);
       // For MMAV2 and V3
       unsigned bitwidth = eltTy.getIntOrFloatBitWidth();
       unsigned kWidth = std::max(32 / bitwidth, 1u);
-      return $_get(context, opIdx, parent, kWidth);
+      return $_get(context, opIdx, parent, kWidth, isChained);
+    }]>,
+
+    AttrBuilder<(ins "unsigned":$opIdx,
+                     "Attribute":$parent,
+                     "Type":$eltTy), [{
+      return get(context, opIdx, parent, eltTy, false);
     }]>
   ];
 

include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h

@@ -150,6 +150,9 @@ Value createAlloc(Operation *insertBefore, RankedTensorType ty, Location loc,
 // Determine if the operation is a TMA load.
 bool isTMALoad(Operation *op);
 
+// Determine if the operation is a PPU AIU load.
+bool isAIULoad(Operation *op);
+
 // Determine if the operation can be lowered to an async load.
 bool canBeAsyncLoad(Operation *op);
 

include/triton/Dialect/TritonGPU/Transforms/Utility.h

@@ -231,6 +231,9 @@ bool isPureUnaryInlineAsm(Operation *op);
 // read the compute capability from the module attributes
 int getNVIDIAComputeCapability(Operation *module);
 
+// read the PPU compute capability from the module attributes
+int getPPUComputeCapability(Operation *module);
+
 // Read the amd target from the module attributes
 std::optional<StringRef> getAMDArch(Operation *module);
 

lib/Analysis/Allocation.cpp

@@ -34,6 +34,7 @@
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "triton/Tools/GenericSwizzling.h"
 #include "triton/Tools/LayoutUtils.h"
@@ -180,7 +181,36 @@ class AllocationAnalysis {
     int64_t bytes =
         numElems * getIntOrFloatOrPtrBitWidth(allocType.getElementType()) / 8;
 
-    auto alignment = alloc.getAlignmentOrDefault();
+    auto mod = op->getParentOfType<ModuleOp>();
+    StringAttr targetAttr =
+        mod->getAttrOfType<StringAttr>(triton::gpu::AttrTargetName);
+    int computeCapability = targetAttr ? getPPUComputeCapability(mod) : 0;
+
+    // Update `bytes` for PPU0015 AIU load
+    if (computeCapability == 89) {
+      if (auto aiuEnc = dyn_cast<triton::gpu::PPUAIUSharedEncodingAttr>(
+              allocType.getEncoding())) {
+        assert(aiuEnc.isPPU0015() && "Unexpected aiuEncoding for PPU0015");
+        auto aiuLoad = aiuEnc.getAIUStrategy();
+        unsigned cubeC = aiuLoad[0];
+        unsigned swizzledBytes = aiuLoad[4];
+        auto elemByteSize = allocType.getElementTypeBitWidth() / 8;
+        if (elemByteSize == 2 && cubeC == 16 ||
+            elemByteSize == 1 && cubeC == 32) {
+          unsigned swizzledElems = swizzledBytes / elemByteSize;
+          unsigned aiuFactor = swizzledElems / cubeC;
+          bytes *= aiuFactor;
+        }
+      }
+    }
+
+    int32_t alignment;
+    if (mlir::isa<triton::gpu::PPUAIUSharedEncodingAttr>(
+            allocType.getEncoding())) {
+      alignment = 128;
+    } else {
+      alignment = alloc.getAlignmentOrDefault();
+    }
     allocation->addBuffer<BufferT::BufferKind::Explicit>(alloc, bytes,
                                                          alignment);
   }