[WIP][LLVM][Flydsl] bump up llvm for upstream and internal coexec opt and adapt flydsl api

.github/workflows/build-custom-llvm-tools.yaml

@@ -86,7 +86,7 @@ jobs:
             REF="$(python3 -c "import json; print(json.load(open('${CONFIG}'))['llvm_ref'])")"
             LLVM_REMOTE="$(python3 -c "import json; print(json.load(open('${CONFIG}')).get('llvm_remote', '${LLVM_REMOTE}'))")"
           else
-            REF="$(tr -d '[:space:]' < flydsl/thirdparty/llvm-hash.txt)"
+            REF="$(python3 -c "import json; print(json.load(open('flydsl/thirdparty/llvm-build-info.json'))['upstream']['llvm_hash'])")"
           fi
 
           # Resolve branch/tag to commit SHA for stable cache key

.github/workflows/build-whl.yaml

@@ -103,7 +103,7 @@ jobs:
         uses: actions/cache@v4
         with:
           path: mlir_install.tgz
-          key: mlir-install-manylinux228-${{ hashFiles('flydsl/thirdparty/llvm-hash.txt', 'flydsl/scripts/build_llvm.sh', 'flydsl/CMakeLists.txt') }}
+          key: mlir-install-manylinux228-${{ hashFiles('flydsl/thirdparty/llvm-build-info.json', 'flydsl/scripts/build_llvm.sh', 'flydsl/CMakeLists.txt') }}
 
       - name: Use cached MLIR install tarball
         if: steps.mlir-cache.outputs.cache-hit == 'true'
@@ -163,7 +163,7 @@ jobs:
           SUMMARY_RELEASE_TYPE: ${{ inputs.release_type }}
           SUMMARY_WHEEL_DIR: dist
         run: |
-          SUMMARY_LLVM_COMMIT="$(tr -d '[:space:]' < flydsl/thirdparty/llvm-hash.txt)"
+          SUMMARY_LLVM_COMMIT="$(python3 -c "import json; print(json.load(open('flydsl/thirdparty/llvm-build-info.json'))['upstream']['llvm_hash'])")"
           export SUMMARY_LLVM_COMMIT
           python3 flydsl/scripts/generate_summary.py build
 

.github/workflows/flydsl.yaml

@@ -147,7 +147,7 @@ jobs:
         uses: actions/cache/restore@v4
         with:
           path: mlir_install.tgz
-          key: mlir-install-${{ runner.os }}-${{ runner.arch }}-${{ env.MLIR_CACHE_VERSION }}-${{ hashFiles('flydsl-test/thirdparty/llvm-hash.txt', 'flydsl-test/scripts/build_llvm.sh') }}
+          key: mlir-install-${{ runner.os }}-${{ runner.arch }}-${{ env.MLIR_CACHE_VERSION }}-${{ hashFiles('flydsl-test/thirdparty/llvm-build-info.json', 'flydsl-test/scripts/build_llvm.sh') }}
 
       - name: Start MLIR build container
         if: steps.mlir-cache.outputs.cache-hit != 'true'

kernels/common/buffer_ops.py

@@ -555,11 +555,11 @@ def buffer_load(
             if not isinstance(soffset.type, ir.IntegerType) or soffset.type.width != 32:
                 op = std_arith.IndexCastOp(T.i32(), soffset)
                 soffset = _unwrap_value(op.result)
-    aux_flags = _create_i32_constant(cache_modifier)
+    aux_attr = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) if cache_modifier else None
 
     # Emit buffer load
     load_op = rocdl.RawPtrBufferLoadOp(
-        result_type, rsrc, offset, soffset, aux_flags  # soffset (scalar byte offset)  # aux (cache modifiers)
+        result_type, rsrc, offset, soffset, aux=aux_attr
     )
 
     return load_op.result
@@ -642,9 +642,9 @@ def buffer_store(
             if not isinstance(soffset.type, ir.IntegerType) or soffset.type.width != 32:
                 op = std_arith.IndexCastOp(T.i32(), soffset)
                 soffset = _unwrap_value(op.result)
-    aux_flags = _create_i32_constant(cache_modifier)
+    aux_attr = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) if cache_modifier else None
 
     # Emit buffer store
     rocdl.RawPtrBufferStoreOp(
-        data, rsrc, offset, soffset, aux_flags  # soffset (scalar byte offset)  # aux (cache modifiers)
+        data, rsrc, offset, soffset, aux=aux_attr
     )

lib/Conversion/FlyToROCDL/FlyToROCDL.cpp

@@ -378,10 +378,11 @@ class PtrLoadOpLowering : public OpConversionPattern<PtrLoadOp> {
     if (isTargetAddressSpace<BufferDescAddressAttr>(flyPtrTy.getAddressSpace())) {
       BufferFatPtr bp(flyPtrTy, ptr);
       Value zero = arith::ConstantIntOp::create(rewriter, loc, 0, 32);
+      auto auxAttr = rewriter.getI32IntegerAttr(0);
       ArrayAttr noAttrs;
       Value loaded = ROCDL::RawPtrBufferLoadOp::create(
           rewriter, loc, loadTy, bp.bufferRsrc(rewriter, loc), bp.swizzleByteOffset(rewriter, loc),
-          zero, zero, noAttrs, noAttrs, noAttrs);
+          zero, auxAttr, noAttrs, noAttrs, noAttrs);
       rewriter.replaceOp(op, loaded);
       return success();
     } else {
@@ -423,10 +424,11 @@ class PtrStoreOpLowering : public OpConversionPattern<PtrStoreOp> {
     if (isTargetAddressSpace<BufferDescAddressAttr>(flyPtrTy.getAddressSpace())) {
       BufferFatPtr bp(flyPtrTy, ptr);
       Value zero = arith::ConstantIntOp::create(rewriter, loc, 0, 32);
+      auto auxAttr = rewriter.getI32IntegerAttr(0);
       ArrayAttr noAttrs;
       ROCDL::RawPtrBufferStoreOp::create(rewriter, loc, value, bp.bufferRsrc(rewriter, loc),
-                                         bp.swizzleByteOffset(rewriter, loc), zero, zero, noAttrs,
-                                         noAttrs, noAttrs);
+                                         bp.swizzleByteOffset(rewriter, loc), zero, auxAttr,
+                                         noAttrs, noAttrs, noAttrs);
       rewriter.eraseOp(op);
       return success();
     } else {

lib/Dialect/Fly/IR/FlyOps.cpp

@@ -190,9 +190,8 @@ Type applyOffsetOnTensorLike(LayoutBuilder<LayoutAttr> &builder, Type tensorLike
 #define FLY_INFER_RETURN_TYPES(OP)                                                                 \
   llvm::LogicalResult OP::inferReturnTypes(                                                        \
       mlir::MLIRContext *context, std::optional<::mlir::Location> location,                        \
-      mlir::ValueRange operands, mlir::DictionaryAttr attributes,                                  \
-      mlir::OpaqueProperties properties, mlir::RegionRange regions,                                \
-      llvm::SmallVectorImpl<mlir::Type> &inferredReturnTypes)
+      mlir::ValueRange operands, mlir::DictionaryAttr attributes, mlir::PropertyRef properties,    \
+      mlir::RegionRange regions, llvm::SmallVectorImpl<mlir::Type> &inferredReturnTypes)
 
 //===----------------------------------------------------------------------===//
 // Constructors

lib/Dialect/Fly/Transforms/PromoteRegMemToVectorSSA.cpp

@@ -27,14 +27,12 @@ using namespace mlir::fly;
 
 namespace llvm {
 
-template <> struct DenseMapInfo<mlir::fly::MakePtrOp> : DenseMapInfo<mlir::Operation *> {
-  using Base = DenseMapInfo<mlir::Operation *>;
-
-  static mlir::fly::MakePtrOp getEmptyKey() { return mlir::fly::MakePtrOp(Base::getEmptyKey()); }
-
-  static mlir::fly::MakePtrOp getTombstoneKey() {
-    return mlir::fly::MakePtrOp(Base::getTombstoneKey());
+template <> struct DenseMapInfo<mlir::fly::MakePtrOp> {
+  static unsigned getHashValue(mlir::fly::MakePtrOp op) {
+    return DenseMapInfo<void *>::getHashValue(op.getAsOpaquePointer());
   }
+
+  static bool isEqual(mlir::fly::MakePtrOp lhs, mlir::fly::MakePtrOp rhs) { return lhs == rhs; }
 };
 
 } // namespace llvm

lib/Dialect/FlyROCDL/CDNA3/CopyAtom.cpp

@@ -89,7 +89,7 @@ FailureOr<Value> CopyOpCDNA3BufferCopyType::emitAtomCallSSA(OpBuilder &builder,
   };
 
   // raw buffer load/store cachepolicy (0=cached, 2=nt)
-  Value aux = arith::ConstantIntOp::create(builder, loc, getCacheModifier(), 32);
+  auto aux = builder.getI32IntegerAttr(getCacheModifier());
   ArrayAttr noAttrs;
 
   auto srcMemTy = srcTyArg ? dyn_cast<fly::MemRefType>(srcTyArg) : fly::MemRefType();
@@ -321,8 +321,9 @@ LogicalResult CopyOpCDNA3BufferCopyLDSType::emitAtomCall(OpBuilder &builder, Loc
   Value srcOff = bp.swizzleByteOffset(builder, loc);
 
   ArrayAttr noAttrs;
+  auto auxAttr = builder.getI32IntegerAttr(0);
   ROCDL::RawPtrBufferLoadLdsOp::create(builder, loc, srcRsrc, dst, size, srcOff, soffset, immOffset,
-                                       zero, noAttrs, noAttrs, noAttrs);
+                                       auxAttr, noAttrs, noAttrs, noAttrs);
   return success();
 }
 
@@ -428,7 +429,7 @@ FailureOr<Value> CopyOpCDNA3BufferAtomicType::emitAtomCallSSA(OpBuilder &builder
     soffset = arith::DivUIOp::create(builder, loc, bits, eight);
   }
 
-  Value zero = arith::ConstantIntOp::create(builder, loc, 0, 32);
+  auto auxAttr = builder.getI32IntegerAttr(0);
   ArrayAttr noAttrs;
 
   AtomicOp op = getAtomicOp().getValue();
@@ -437,22 +438,22 @@ FailureOr<Value> CopyOpCDNA3BufferAtomicType::emitAtomCallSSA(OpBuilder &builder
   case AtomicOp::Add:
     if (!isFloat)
       return failure();
-    ROCDL::RawPtrBufferAtomicFaddOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero,
-                                            noAttrs, noAttrs, noAttrs);
+    ROCDL::RawPtrBufferAtomicFaddOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff,
+                                            soffset, auxAttr, noAttrs, noAttrs, noAttrs);
     break;
   case AtomicOp::Max:
     if (isFloat)
-      ROCDL::RawPtrBufferAtomicFmaxOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero,
-                                              noAttrs, noAttrs, noAttrs);
+      ROCDL::RawPtrBufferAtomicFmaxOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff,
+                                              soffset, auxAttr, noAttrs, noAttrs, noAttrs);
     else
-      ROCDL::RawPtrBufferAtomicSmaxOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero,
-                                              noAttrs, noAttrs, noAttrs);
+      ROCDL::RawPtrBufferAtomicSmaxOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff,
+                                              soffset, auxAttr, noAttrs, noAttrs, noAttrs);
     break;
   case AtomicOp::Min:
     if (isFloat)
       return failure();
-    ROCDL::RawPtrBufferAtomicUminOp::create(builder, loc, src, dstRsrc, dstOff, soffset, zero,
-                                            noAttrs, noAttrs, noAttrs);
+    ROCDL::RawPtrBufferAtomicUminOp::create(builder, loc, src.getType(), src, dstRsrc, dstOff,
+                                            soffset, auxAttr, noAttrs, noAttrs, noAttrs);
     break;
   default:
     return failure();

lib/Dialect/FlyROCDL/CDNA3/MmaAtom.cpp

@@ -169,8 +169,8 @@ FailureOr<Value> MmaOpCDNA3_MFMAType::emitAtomCallSSA(OpBuilder &builder, Locati
 
 #define DISPATCH_MFMA_SSA(M_, K_, PRED, OP)                                                        \
   if (m == M_ && n == M_ && k == K_ && (PRED)) {                                                   \
-    auto zeroAttr = builder.getI32IntegerAttr(0);                                                  \
-    return ROCDL::OP::create(builder, loc, accTy, a, b, c, zeroAttr, zeroAttr, zeroAttr)           \
+    return ROCDL::OP::create(builder, loc, accTy, a, b, c, (uint32_t)0, (uint32_t)0,               \
+                             ROCDL::MFMAPermB::none)                                               \
         .getResult();                                                                              \
   }
 

lib/Dialect/FlyROCDL/CDNA4/MmaAtom.cpp

@@ -237,21 +237,19 @@ FailureOr<Value> MmaOpCDNA4_MFMAScaleType::emitAtomCallSSA(OpBuilder &builder, L
   Value scaleB = LLVM::ExtractValueOp::create(
       builder, loc, atomVal, ArrayRef<int64_t>{*getFieldIndex(AtomStateField::ScaleB)});
 
-  auto cbszAttr = builder.getI32IntegerAttr(*aTypeCode);
-  auto blgpAttr = builder.getI32IntegerAttr(*bTypeCode);
-  auto opselAAttr = builder.getI32IntegerAttr(getOpselA());
-  auto opselBAttr = builder.getI32IntegerAttr(getOpselB());
+  auto cbsz = static_cast<ROCDL::MatrixFormat>(*aTypeCode);
+  auto blgp = static_cast<ROCDL::MatrixFormat>(*bTypeCode);
+  uint32_t opselA = getOpselA();
+  uint32_t opselB = getOpselB();
 
   if (m == 16 && n == 16 && k == 128) {
-    return ROCDL::mfma_scale_f32_16x16x128_f8f6f4::create(builder, loc, accTy, a, b, c, cbszAttr,
-                                                          blgpAttr, opselAAttr, scaleA, opselBAttr,
-                                                          scaleB)
+    return ROCDL::mfma_scale_f32_16x16x128_f8f6f4::create(builder, loc, accTy, a, b, c, cbsz, blgp,
+                                                          opselA, scaleA, opselB, scaleB)
         .getResult();
   }
   if (m == 32 && n == 32 && k == 64) {
-    return ROCDL::mfma_scale_f32_32x32x64_f8f6f4::create(builder, loc, accTy, a, b, c, cbszAttr,
-                                                         blgpAttr, opselAAttr, scaleA, opselBAttr,
-                                                         scaleB)
+    return ROCDL::mfma_scale_f32_32x32x64_f8f6f4::create(builder, loc, accTy, a, b, c, cbsz, blgp,
+                                                         opselA, scaleA, opselB, scaleB)
         .getResult();
   }
 

lib/Dialect/FlyROCDL/GFX1250/CopyAtom.cpp

@@ -503,7 +503,7 @@ LogicalResult CopyOpGFX1250TDMType::emitAtomCall(OpBuilder &builder, Location lo
       vector::FromElementsOp::create(builder, loc, VectorType::get({8}, builder.getI32Type()),
                                      ValueRange{zero, zero, zero, zero, zero, zero, zero, zero});
 
-  uint32_t cachePolicy = static_cast<uint32_t>(getCacheModifier());
+  auto cachePolicy = builder.getI32IntegerAttr(static_cast<int32_t>(getCacheModifier()));
   ArrayAttr noAliasScopes;
   if (isLoad)
     ROCDL::TensorLoadToLDSOp::create(builder, loc, dgroup0, dgroup1, dg2, dg3, dg4, cachePolicy,

lib/Dialect/FlyROCDL/GFX1250/MmaAtom.cpp

@@ -238,24 +238,18 @@ static FailureOr<Value> emitWmmaSSA(OpBuilder &builder, Location loc, VectorType
                                     bool clamp = false) {
   Value res;
   if constexpr (Variant == WmmaVariant::ModsAllReuse) {
-    // Float path: no sign/clamp operands.
-    res = WmmaOp::create(builder, loc, accTy,
-                         /*signA=*/false, a, /*signB=*/false, b,
-                         /*modC=*/(uint16_t)0, c)
+    res = WmmaOp::create(builder, loc, accTy, a, b, ROCDL::WMMACModifier::none, c,
+                         /*reuseA=*/false, /*reuseB=*/false)
               .getResult();
   } else if constexpr (Variant == WmmaVariant::ModsC) {
-    // fp8 path: no sign/clamp operands.
-    res = WmmaOp::create(builder, loc, accTy, a, b,
-                         /*modC=*/(uint16_t)0, c,
+    res = WmmaOp::create(builder, loc, accTy, a, b, ROCDL::WMMACModifier::none, c,
                          /*reuseA=*/false, /*reuseB=*/false)
               .getResult();
   } else if constexpr (Variant == WmmaVariant::ModsABClamp) {
-    // iu8: sign + reuse + clamp controls.
     res = WmmaOp::create(builder, loc, accTy, signA, a, signB, b, c,
                          /*reuseA=*/false, /*reuseB=*/false, clamp)
               .getResult();
   } else {
-    // IU form (e.g. iu4): sign/clamp controls but no reuseA/reuseB operands.
     static_assert(Variant == WmmaVariant::ModsIUClamp);
     res = WmmaOp::create(builder, loc, accTy, signA, a, signB, b, c, clamp).getResult();
   }

lib/Dialect/FlyROCDL/GFX1250/MmaAtomScale.cpp

@@ -240,20 +240,20 @@ FailureOr<Value> MmaOpGFX1250_WMMAScaleType::emitAtomCallSSA(OpBuilder &builder,
   // the atom's compile-time params. block-16 selects the V_WMMA_SCALE16 form
   // (i64 scale operands); block-32 the V_WMMA_SCALE form (i32 scale operands).
   bool block16 = getBlockSize() == 16;
+  auto modC = static_cast<ROCDL::WMMACModifier>(getModC());
+  auto scaleAType = static_cast<ROCDL::WMMAMatrixScale>(getOpselA());
+  auto scaleBType = static_cast<ROCDL::WMMAMatrixScale>(getOpselB());
+  auto fmtScale0 = static_cast<ROCDL::WMMAMatrixScaleFormat>(0);
+
   if (m == 32 && n == 16 && k == 128) {
-    // fp4-only form; no fmtA/fmtB operands.
     if (block16)
       return ROCDL::wmma_scale16_f32_32x16x128_f4::create(
-                 builder, loc, accTy, a, b, /*modC=*/(uint16_t)getModC(), c,
-                 /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA,
-                 /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB,
-                 /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB())
+                 builder, loc, accTy, a, b, modC, c, scaleAType, fmtScale0, scaleA, scaleBType,
+                 fmtScale0, scaleB, getReuseA(), getReuseB())
           .getResult();
-    return ROCDL::wmma_scale_f32_32x16x128_f4::create(
-               builder, loc, accTy, a, b, /*modC=*/(uint16_t)getModC(), c,
-               /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA,
-               /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB,
-               /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB())
+    return ROCDL::wmma_scale_f32_32x16x128_f4::create(builder, loc, accTy, a, b, modC, c,
+                                                      scaleAType, fmtScale0, scaleA, scaleBType,
+                                                      fmtScale0, scaleB, getReuseA(), getReuseB())
         .getResult();
   }
 
@@ -262,20 +262,17 @@ FailureOr<Value> MmaOpGFX1250_WMMAScaleType::emitAtomCallSSA(OpBuilder &builder,
   if (!aFmt || !bFmt)
     return failure();
 
+  auto fmtA = static_cast<ROCDL::MatrixFormat>(*aFmt);
+  auto fmtB = static_cast<ROCDL::MatrixFormat>(*bFmt);
+
   if (block16)
     return ROCDL::wmma_scale16_f32_16x16x128_f8f6f4::create(
-               builder, loc, accTy, /*fmtA=*/*aFmt, a, /*fmtB=*/*bFmt, b,
-               /*modC=*/(uint16_t)getModC(), c,
-               /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA,
-               /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB,
-               /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB())
+               builder, loc, accTy, fmtA, a, fmtB, b, modC, c, scaleAType, fmtScale0, scaleA,
+               scaleBType, fmtScale0, scaleB, getReuseA(), getReuseB())
         .getResult();
   return ROCDL::wmma_scale_f32_16x16x128_f8f6f4::create(
-             builder, loc, accTy, /*fmtA=*/*aFmt, a, /*fmtB=*/*bFmt, b,
-             /*modC=*/(uint16_t)getModC(), c,
-             /*scaleAType=*/(uint32_t)getOpselA(), /*fmtScaleA=*/(uint32_t)0, scaleA,
-             /*scaleBType=*/(uint32_t)getOpselB(), /*fmtScaleB=*/(uint32_t)0, scaleB,
-             /*reuseA=*/getReuseA(), /*reuseB=*/getReuseB())
+             builder, loc, accTy, fmtA, a, fmtB, b, modC, c, scaleAType, fmtScale0, scaleA,
+             scaleBType, fmtScale0, scaleB, getReuseA(), getReuseB())
       .getResult();
 }
 

lib/Dialect/FlyROCDL/Ops.cpp

@@ -15,7 +15,7 @@ using namespace mlir::fly_rocdl;
 LogicalResult GetBufferRsrcOp::inferReturnTypes(MLIRContext *context,
                                                 std::optional<Location> location,
                                                 ValueRange operands, DictionaryAttr attributes,
-                                                OpaqueProperties properties, RegionRange regions,
+                                                PropertyRef properties, RegionRange regions,
                                                 SmallVectorImpl<Type> &inferredReturnTypes) {
   auto ptrTy = dyn_cast<PointerType>(operands[0].getType());
   if (!ptrTy)