From bab73389658060692205d50a47a27aa9009a4095 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Tue, 6 Dec 2022 13:24:50 -0800 Subject: [PATCH] Fix --- .../TritonGPUToLLVM/TritonGPUToLLVM.cpp | 6 ++- test/Conversion/tritongpu_to_llvm.mlir | 51 +++---------------- 2 files changed, 10 insertions(+), 47 deletions(-) diff --git a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp index a5fb4edc0..b0ee3336d 100644 --- a/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp +++ b/lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp @@ -4105,10 +4105,12 @@ struct InsertSliceAsyncOpConversion } v = bitcast(v, IntegerType::get(getContext(), byteWidth)); // Write shared memory if predicate is true - auto *valOperand = ptxBuilder.newOperand(v, "r"); - auto &st = *ptxBuilder.create("st"); + PTXBuilder ptxStoreBuilder; + auto *valOperand = ptxStoreBuilder.newOperand(v, "r"); + auto &st = *ptxStoreBuilder.create("st"); st.shared().o("b" + std::to_string(bitWidth)); st(dstOperand, valOperand).predicate(pred); + ptxStoreBuilder.launch(rewriter, loc, void_ty(getContext())); } } copyAsyncOp(dstOperand, srcOperand, copySize, srcSize); diff --git a/test/Conversion/tritongpu_to_llvm.mlir b/test/Conversion/tritongpu_to_llvm.mlir index c897dde4d..3135ee490 100644 --- a/test/Conversion/tritongpu_to_llvm.mlir +++ b/test/Conversion/tritongpu_to_llvm.mlir @@ -504,11 +504,13 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { %other_tensor = tt.splat %other : (f32) -> tensor<16x64xf32, #AL> // CHECK: llvm.inline_asm has_side_effects asm_dialect = att - // CHECK-SAME: st.shared.b128 [ ${{.*}} + 0 ] - // CHECK-SAME: cp.async.cg.shared.global [ ${{.*}} + 0 ], [ ${{.*}} + 0 ], 0x10, 0x10 + // CHECK-SAME: @${{.*}} st.shared.b128 [ ${{.*}} + 0 ], ${{.*}} // CHECK: llvm.inline_asm has_side_effects asm_dialect = att - // CHECK-SAME: st.shared.b128 [ ${{.*}} + 16 ] - // CHECK-SAME: cp.async.cg.shared.global [ ${{.*}} + 16 ], [ ${{.*}} + 0 ], 0x10, 0x10 + // CHECK-SAME: cp.async.cg.shared.global [ ${{.*}} + 0 ], [ ${{.*}} + 0 ], 0x10, ${{.*}} + // CHECK: llvm.inline_asm has_side_effects asm_dialect = att + // CHECK-SAME: @${{.*}} st.shared.b128 [ ${{.*}} + 16 ], ${{.*}} + // CHECK: llvm.inline_asm has_side_effects asm_dialect = att + // CHECK-SAME: cp.async.cg.shared.global [ ${{.*}} + 16 ], [ ${{.*}} + 0 ], 0x10, ${{.*}} // CHECK: llvm.inline_asm has_side_effects asm_dialect = att // CHECK-SAME: cp.async.commit_group %a = triton_gpu.insert_slice_async %a_ptr, %tensor, %index, %true_tensor, %other_tensor {axis = 0 : i32, cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<16x64x!tt.ptr, #AL> -> tensor<2x16x64xf32, #A> @@ -518,47 +520,6 @@ module attributes {"triton_gpu.num-warps" = 4 : i32} { // ----- -#block0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [4], warpsPerCTA = [4], order = [0]}> -#block1 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [8], warpsPerCTA = [4], order = [0]}> -#block2 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [4, 1], warpsPerCTA = [4, 1], order = [1, 0]}> -#block3 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 8], warpsPerCTA = [1, 4], order = [1, 0]}> -#slice2d1 = #triton_gpu.slice<{dim = 1, parent=#block2}> -#slice3d0 = #triton_gpu.slice<{dim = 0, parent=#block3}> -#AL = #triton_gpu.blocked<{sizePerThread = [1, 8], threadsPerWarp = [4, 8], warpsPerCTA = [4, 1], order = [1, 0]}> -#A = #triton_gpu.shared<{vec = 8, perPhase = 1, maxPhase = 4, order = [1, 0]}> -module attributes {"triton_gpu.num-warps" = 4 : i32} { - // CHECK-LABEL: basic_insert_slice_async_v4 - func @basic_insert_slice_async_v4(%arg0: !tt.ptr {tt.divisibility = 8 : i32}) { - %off0_ = tt.make_range {end = 16 : i32, start = 0 : i32} : tensor<16xi32, #slice2d1> - %off1_ = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<64xi32, #slice3d0> - %off0 = tt.expand_dims %off0_ {axis = 1 : i32} : (tensor<16xi32, #slice2d1>) -> tensor<16x1xi32, #block2> - %off1 = tt.expand_dims %off1_ {axis = 0 : i32} : (tensor<64xi32, #slice3d0>) -> tensor<1x64xi32, #block3> - %broadcast_off0_scalar = tt.broadcast %off0 : (tensor<16x1xi32, #block2>) -> tensor<16x64xi32, #block2> - %cst_scalar = arith.constant 64 : i32 - %cst = tt.splat %cst_scalar : (i32) -> tensor<16x64xi32, #block2> - %broadcast_off0_ = arith.muli %broadcast_off0_scalar, %cst : tensor<16x64xi32, #block2> - %broadcast_off1_ = tt.broadcast %off1 : (tensor<1x64xi32, #block3>) -> tensor<16x64xi32, #block3> - %broadcast_off0 = triton_gpu.convert_layout %broadcast_off0_ : (tensor<16x64xi32, #block2>) -> tensor<16x64xi32, #AL> - %broadcast_off1 = triton_gpu.convert_layout %broadcast_off1_ : (tensor<16x64xi32, #block3>) -> tensor<16x64xi32, #AL> - %off = arith.addi %broadcast_off0, %broadcast_off1 : tensor<16x64xi32, #AL> - %a_init = tt.splat %arg0 : (!tt.ptr) -> tensor<16x64x!tt.ptr, #AL> - %a_ptr = tt.addptr %a_init, %off : tensor<16x64x!tt.ptr, #AL> - %tensor = triton_gpu.alloc_tensor : tensor<2x16x64xf32, #A> - %index = arith.constant 1 : i32 - - // CHECK: llvm.inline_asm has_side_effects asm_dialect = att - // CHECK-SAME: cp.async.cg.shared.global [ ${{.*}} + 0 ], [ ${{.*}} + 0 ], 0x10, 0x10 - // CHECK: llvm.inline_asm has_side_effects asm_dialect = att - // CHECK-SAME: cp.async.cg.shared.global [ ${{.*}} + 16 ], [ ${{.*}} + 0 ], 0x10, 0x10 - // CHECK: llvm.inline_asm has_side_effects asm_dialect = att - // CHECK-SAME: cp.async.commit_group - %a = triton_gpu.insert_slice_async %a_ptr, %tensor, %index {axis = 0 : i32, cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<16x64x!tt.ptr, #AL> -> tensor<2x16x64xf32, #A> - return - } -} - -// ----- - #block0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [4], warpsPerCTA = [4], order = [0]}> #block1 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [8], warpsPerCTA = [4], order = [0]}> #block2 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [4, 1], warpsPerCTA = [4, 1], order = [1, 0]}>