Merge commit 'b5fea1e3f4c2cb0b40c0ce98261b240d8728d2f9'

Author: whitneywhtsang

.github/workflows/integration-tests-amd.yml

@@ -181,12 +181,7 @@ jobs:
       - name: Run Proton tests
         run: |
           unset HIP_VISIBLE_DEVICES
-          unset ROCR_VISIBLE_DEVICES
-          if [ "${{ matrix.runner[0] }}" = "amd-gfx950" ]; then
-            python3 -m pytest -s -n 8 third_party/proton/test -k "not test_instrument_exec"
-          else
-            make test-proton
-          fi
+          make test-proton
       - name: Inspect cache directories
         run: |
           mkdir -p ~/.triton

.github/workflows/llvm-build.yml

@@ -31,7 +31,6 @@ jobs:
         config:
         - {runner: 'Ubuntu 22.04', runs_on: 'ubuntu-22.04', target-os: 'ubuntu', arch: 'x64'}
         - {runner: 'Ubuntu 22.04 ARM64', runs_on: 'ubuntu-22.04', target-os: 'ubuntu', arch: 'arm64'}
-        - {runner: 'CentOS 7', runs_on: ['self-hosted', 'CPU'], target-os: 'centos', arch: 'x64'}
         - {runner: 'AlmaLinux 8', runs_on: ['self-hosted', 'CPU'], target-os: 'almalinux', arch: 'x64'}
         - {runner: 'AlmaLinux 8 ARM64', runs_on: 'ubuntu-22.04-arm', target-os: 'almalinux', arch: 'arm64'}
         - {runner: 'MacOS X64', runs_on: 'macos-13', target-os: 'macos', arch: 'x64'}
@@ -233,30 +232,6 @@ jobs:
 
         tar czf "${{ env.llvm_install_dir }}.tar.gz" "${{ env.llvm_install_dir }}"
 
-
-    - name: Configure, Build, Test, and Install LLVM (CentOS)
-      if: matrix.config.target-os == 'centos'
-      run: |
-        # if this step crashes, it can leave behind a stale docker container
-        docker container prune -f
-        docker rmi -f $(docker images -q)
-
-        docker build --tag llvm-build --build-arg llvm_dir=llvm-project \
-          -f llvm-build/.github/workflows/llvm-build/centos.Dockerfile .
-
-        # Create temporary container to copy cache and installed artifacts.
-        CONTAINER_ID=$(docker create llvm-build)
-        docker cp "${CONTAINER_ID}:/install" "${{ env.llvm_install_dir }}"
-        tar czf "${{ env.llvm_install_dir }}.tar.gz" "${{ env.llvm_install_dir }}"
-
-        # We remove the existing directory, otherwise docker will
-        # create a subdirectory inside the existing directory.
-        rm -rf "${{ env.SCCACHE_DIR }}"
-        docker cp "${CONTAINER_ID}:/sccache" "${{ env.SCCACHE_DIR }}"
-        sudo chown -R "$(id -u -n):$(id -g -n)" "${{ env.SCCACHE_DIR }}"
-
-        docker rm "${CONTAINER_ID}"
-
     - name: Configure, Build, Test, and Install LLVM (AlmaLinux)
       if: matrix.config.target-os == 'almalinux'
       run: |

.github/workflows/llvm-build/centos.Dockerfile

@@ -6,8 +6,6 @@ if(POLICY CMP0116)
   cmake_policy(SET CMP0116 OLD)
 endif()
 
-include(ExternalProject)
-
 set(CMAKE_CXX_STANDARD 17)
 
 set(CMAKE_INCLUDE_CURRENT_DIR ON)

CMakeLists.txt

@@ -73,14 +73,91 @@ struct ReturnOpConversion : public OpConversionPattern<ReturnOp> {
   }
 };
 
+//===----------------------------------------------------------------------===//
+// FunctionOpInterfaceSignatureConversion
+//===----------------------------------------------------------------------===//
+// NOTE: Forked from mlir to support remapping argument attributes correctly in
+// a one-to-many type conversion.
+
+SmallVector<Attribute>
+convertFuncOpAttrs(FunctionOpInterface funcOp,
+                   TypeConverter::SignatureConversion &sigConv,
+                   FunctionType newType) {
+  if (newType.getNumInputs() == funcOp.getNumArguments()) {
+    return {};
+  }
+  ArrayAttr allArgAttrs = funcOp.getAllArgAttrs();
+  if (!allArgAttrs)
+    return {};
+
+  SmallVector<Attribute> newAttrs(newType.getNumInputs());
+  for (auto i : llvm::seq(allArgAttrs.size())) {
+    auto mapping = sigConv.getInputMapping(i);
+    assert(mapping.has_value());
+    auto outIdx = mapping->inputNo;
+    newAttrs[outIdx] = allArgAttrs[i];
+  }
+  return newAttrs;
+}
+
+LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
+                                 const TypeConverter &typeConverter,
+                                 ConversionPatternRewriter &rewriter) {
+  FunctionType type = dyn_cast<FunctionType>(funcOp.getFunctionType());
+  if (!type)
+    return failure();
+
+  // Convert the original function types.
+  TypeConverter::SignatureConversion result(type.getNumInputs());
+  SmallVector<Type, 1> newResults;
+  if (failed(typeConverter.convertSignatureArgs(type.getInputs(), result)) ||
+      failed(typeConverter.convertTypes(type.getResults(), newResults)) ||
+      failed(rewriter.convertRegionTypes(&funcOp.getFunctionBody(),
+                                         typeConverter, &result)))
+    return failure();
+
+  // Update the function signature in-place.
+  auto newType = FunctionType::get(rewriter.getContext(),
+                                   result.getConvertedTypes(), newResults);
+
+  auto newArgAttrs = convertFuncOpAttrs(funcOp, result, newType);
+
+  rewriter.modifyOpInPlace(funcOp, [&] {
+    funcOp.setType(newType);
+    if (!newArgAttrs.empty()) {
+      funcOp.setAllArgAttrs(newArgAttrs);
+    }
+  });
+
+  return success();
+}
+
+/// Create a default conversion pattern that rewrites the type signature of a
+/// FunctionOpInterface op. This only supports ops which use FunctionType to
+/// represent their type.
+struct FunctionOpInterfaceSignatureConversion : public ConversionPattern {
+  FunctionOpInterfaceSignatureConversion(StringRef functionLikeOpName,
+                                         MLIRContext *ctx,
+                                         const TypeConverter &converter,
+                                         PatternBenefit benefit = 1)
+      : ConversionPattern(converter, functionLikeOpName, benefit, ctx) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> /*operands*/,
+                  ConversionPatternRewriter &rewriter) const override {
+    FunctionOpInterface funcOp = cast<FunctionOpInterface>(op);
+    return convertFuncOpTypes(funcOp, *typeConverter, rewriter);
+  }
+};
+
 } // namespace
 
 void populateFunctionTypeConversions(const TypeConverter &converter,
                                      RewritePatternSet &patterns) {
-  mlir::populateFunctionOpInterfaceTypeConversionPattern<mlir::triton::FuncOp>(
-      patterns, converter);
-  patterns.add<CallOpConversion, ReturnOpConversion>(converter,
-                                                     patterns.getContext());
+  auto context = patterns.getContext();
+  patterns.add<FunctionOpInterfaceSignatureConversion>(
+      triton::FuncOp::getOperationName(), context, converter);
+  patterns.add<CallOpConversion, ReturnOpConversion>(converter, context);
 }
 
 } // namespace mlir::triton

lib/Dialect/Triton/Transforms/FunctionTypeConversion.cpp

@@ -754,17 +754,19 @@ void createBarrierAndWaitOps(scf::ForOp forOp, CoarseSchedule &schedule,
   Value barrierAlloc = createBarrierAlloc(forOp, numStages);
   Value vTrue = builder.create<arith::ConstantIntOp>(1, 1);
   Value phase = forOp.getRegionIterArg(phaseArgIdx);
-  Value barrierIdx = forOp.getRegionIterArg(barrierIdxArgIdx);
   Value zero = builder.create<arith::ConstantIntOp>(forOp.getLoc(), 0, 32);
+  Value barrierIdx;
+  if (numStages > 1) {
+    barrierIdx = forOp.getRegionIterArg(barrierIdxArgIdx);
+  } else {
+    barrierIdx = zero;
+  }
   Value one = builder.create<arith::ConstantIntOp>(forOp.getLoc(), 1, 32);
   Value numStagesVal =
       builder.create<arith::ConstantIntOp>(forOp.getLoc(), numStages, 32);
 
-  Value barrierSlice = barrierAlloc;
-  if (numStages > 1) {
-    barrierSlice =
-        triton::createSingleBufferView(builder, barrierAlloc, barrierIdx);
-  }
+  Value barrierSlice =
+      triton::createSingleBufferView(builder, barrierAlloc, barrierIdx);
   mma.addCompletionBarrier(barrierSlice, vTrue);
   mma.setIsAsync(true);
 

lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp

@@ -812,6 +812,12 @@ void init_gluon_ir(py::module &&m) {
              self.create<ttag::AsyncTDMCopyGlobalToLocalOp>(descPtr, indices,
                                                             result, pred);
            })
+      .def("create_async_tdm_copy_local_to_global",
+           [](GluonOpBuilder &self, Value descPtr, std::vector<Value> &indices,
+              Value src) {
+             self.create<ttag::AsyncTDMCopyLocalToGlobalOp>(descPtr, indices,
+                                                            src);
+           })
       .def("create_async_tdm_wait", [](GluonOpBuilder &self, int num) {
         ValueRange tokens;
         self.create<ttag::AsyncTDMWait>(tokens, num);

python/src/gluon_ir.cc

@@ -30,7 +30,9 @@
     TensorMemoryScalesLayout,
     allocate_tensor_memory,
     get_tmem_32x32b_reg_layout,
+    get_tmem_scales_reg_layout,
     tcgen05_mma,
+    tcgen05_mma_scaled,
     tcgen05_commit,
     tcgen05_copy,
     float2,
@@ -1334,3 +1336,92 @@ def kernel_auto_layout_constant(threads_per_warp: ttgl.constexpr):
 
 def test_auto_layout_constant():
     kernel_auto_layout_constant.warmup(THREADS_PER_WARP, grid=(1, ))
+
+
+def fp8e8m0_to_float32(scale):
+    scale = scale.view(torch.uint8)
+    scale = scale.to(torch.int32)
+    scale = scale << 23
+    scale = scale.view(torch.float32)
+    return scale
+
+
+@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
+def test_tcgen05_mma_scaled_minimal():
+    M = 128
+    N = 128
+    K = 128
+    threads_per_warp = ttgl.constexpr(THREADS_PER_WARP)
+
+    @gluon.jit
+    def kernel(out_ptr, M: ttgl.constexpr, N: ttgl.constexpr, K: ttgl.constexpr, a, b, a_scale, b_scale):
+        # Simple register layout for creating constants and storing results
+        reg_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [threads_per_warp, 1], [ttgl.num_warps(), 1], [1, 0])
+
+        # Shared-memory layouts for MMA operands
+        nvmma_layout: ttgl.constexpr = ttgl.NVMMASharedLayout(swizzle_byte_width=128, transposed=False,
+                                                              element_bitwidth=8, rank=2)
+        # Allocate zero operands in shared memory (values don't matter since scales are zero)
+        block_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1], [1, 32], warps_per_cta=[ttgl.num_warps(), 1],
+                                                          order=[1, 0])
+        a_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, block_layout))[:, None]
+        a_offs_k = ttgl.arange(0, K, layout=ttgl.SliceLayout(0, block_layout))[None, :]
+        b_offs_k = ttgl.arange(0, K, layout=ttgl.SliceLayout(1, block_layout))[:, None]
+        b_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, block_layout))[None, :]
+
+        a_tile = ttgl.load(a + a_offs_m * K + a_offs_k)
+        b_tile = ttgl.load(b + b_offs_k * N + b_offs_n)
+        a_smem = ttgl.allocate_shared_memory(ttgl.float8e5, [M, K], nvmma_layout, a_tile)
+        b_smem = ttgl.allocate_shared_memory(ttgl.float8e5, [K, N], nvmma_layout, b_tile)
+
+        # Accumulator in TMEM initialized to ones
+        acc_tmem_layout: ttgl.constexpr = TensorMemoryLayout([M, N], col_stride=1)
+        tmem_reg_layout: ttgl.constexpr = get_tmem_32x32b_reg_layout(M, N, [M, N], ttgl.num_warps())
+        acc_init = ttgl.zeros([M, N], ttgl.float32, layout=tmem_reg_layout)
+        acc_tmem = allocate_tensor_memory(ttgl.float32, [M, N], acc_tmem_layout, acc_init)
+
+        # Zero scales in TMEM
+        scale_layout: ttgl.constexpr = TensorMemoryScalesLayout()
+        scale_reg_layout: ttgl.constexpr = get_tmem_scales_reg_layout(M, N, [M, N], ttgl.num_warps())
+        scale_offs_k = ttgl.arange(0, (K // 32), layout=ttgl.SliceLayout(0, scale_reg_layout))[None, :]
+        scale_offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, scale_reg_layout))[:, None]
+        scale_offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(1, scale_reg_layout))[:, None]
+        a_scale_init = ttgl.load(a_scale + scale_offs_m * (K // 32) + scale_offs_k)
+        b_scale_init = ttgl.load(b_scale + scale_offs_n * (K // 32) + scale_offs_k)
+        a_scale_tmem = allocate_tensor_memory(ttgl.int8, [M, K // 32], scale_layout, a_scale_init)
+        b_scale_tmem = allocate_tensor_memory(ttgl.int8, [M, K // 32], scale_layout, b_scale_init)
+
+        # Issue a single scaled MMA and commit
+        bar = ttgl.allocate_shared_memory(ttgl.int64, [1], mbarrier.MBarrierLayout())
+        mbarrier.init(bar, count=1)
+        tcgen05_mma_scaled(a_smem, b_smem, acc_tmem, a_scale_tmem, b_scale_tmem, "e5m2", "e5m2", use_acc=True)
+        tcgen05_commit(bar)
+        mbarrier.wait(bar, phase=0)
+
+        # Load result from TMEM and store to global
+        out_reg = acc_tmem.load(tmem_reg_layout)
+        store_layout: ttgl.constexpr = reg_layout
+        offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, store_layout))[:, None]
+        offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, store_layout))[None, :]
+        offs = offs_m * N + offs_n
+        ttgl.store(out_ptr + offs, ttgl.convert_layout(out_reg, store_layout))
+
+    out = torch.empty((M, N), dtype=torch.float32, device="cuda")
+    a = torch.randint(20, 40, (M, K), dtype=torch.uint8, device="cuda").view(torch.float8_e5m2)
+    b = torch.randint(20, 40, (K, N), dtype=torch.uint8, device="cuda").view(torch.float8_e5m2)
+    a_scale = torch.randint(64, 130, (M, K // 32), dtype=torch.uint8, device="cuda")
+    b_scale = torch.randint(64, 130, (N, K // 32), dtype=torch.uint8, device="cuda")
+    compiled = kernel[(1, )](out, M, N, K, a, b, a_scale, b_scale)
+    A = a.to(torch.float32)
+    B = b.to(torch.float32)
+    a_scale_f32 = fp8e8m0_to_float32(a_scale)
+    b_scale_f32 = fp8e8m0_to_float32(b_scale)
+    a_scale_f32 = a_scale_f32.repeat_interleave(32, dim=1)
+    b_scale_f32 = b_scale_f32.repeat_interleave(32, dim=1)
+    b_scale_f32 = b_scale_f32.T.contiguous()
+    A = A * a_scale_f32
+    B = B * b_scale_f32
+    ref = torch.matmul(A, B)
+    torch.testing.assert_close(out, ref, atol=1e-6, rtol=1e-6)
+    ttgir = compiled.asm["ttgir"]
+    assert "ttng.tc_gen5_mma_scaled" in ttgir