[WS] make aref ops stage/phase optional (triton-lang#8101)

* make aref ops stage/phase optional
* update lit tests

Author: 3gx

lib/Dialect/TritonGPU/Transforms/WarpSpecialization/RewritePartitionDependencies.cpp

@@ -64,16 +64,14 @@ namespace {
 struct AsyncRef {
   auto putView(PartitionBuilder &b, Partition &partition,
                StageCluster srcStageCluster) {
-    auto zero = b.create<arith::ConstantOp>(b.getI32IntegerAttr(0));
     auto enterOp = b.createInto<triton::nvws::ArefPutEnterOp>(
-        partition, srcStageCluster, viewType, tokenType, aref, zero, zero);
+        partition, srcStageCluster, aref, TypeRange{viewType}, tokenType);
     auto token = enterOp.getToken();
 
     auto exitOp = [this, &partition, srcStageCluster,
                    token](PartitionBuilder &b) {
-      auto zero = b.create<arith::ConstantOp>(b.getI32IntegerAttr(0));
       auto exitOp = b.createInto<triton::nvws::ArefPutExitOp>(
-          partition, srcStageCluster, aref, token, zero,
+          partition, srcStageCluster, aref, token,
           b.getArrayAttr(SmallVector<Attribute>{triton::nvws::AsyncOpAttr::get(
               aref.getContext(), triton::nvws::AsyncOp::NONE)}));
     };
@@ -82,16 +80,14 @@ struct AsyncRef {
 
   auto getView(PartitionBuilder &b, Partition &partition,
                StageCluster srcStageCluster) {
-    auto zero = b.create<arith::ConstantOp>(b.getI32IntegerAttr(0));
     auto enterOp = b.createInto<triton::nvws::ArefGetEnterOp>(
-        partition, srcStageCluster, viewType, tokenType, aref, zero, zero);
+        partition, srcStageCluster, aref, TypeRange{viewType}, tokenType);
     auto token = enterOp.getToken();
 
     auto exitOp = [this, &partition, srcStageCluster,
                    token](PartitionBuilder &b) {
-      auto zero = b.create<arith::ConstantOp>(b.getI32IntegerAttr(0));
       auto exitOp = b.createInto<triton::nvws::ArefGetExitOp>(
-          partition, srcStageCluster, aref, token, zero,
+          partition, srcStageCluster, aref, token,
           b.getArrayAttr(SmallVector<Attribute>{triton::nvws::AsyncOpAttr::get(
               aref.getContext(), triton::nvws::AsyncOp::NONE)}));
     };

test/NVWS/insert_aref.mlir

@@ -30,13 +30,10 @@ module attributes {"ttg.num-warps" = 4 : i32, ttg.target = "cuda:100"} {
     // CHECK-NEXT: [[AREF2:%.*]] = nvws.aref.create [[AREF_BUF2]]
     %1 = scf.for %arg5 = %c0_i32 to %arg0 step %c1_i32 iter_args(%arg6 = %0) -> (!ttg.async.token)  : i32 {
       %2 = arith.muli %arg5, %c64_i32 {loop.cluster = 1 : i32, loop.stage = 0 : i32} : i32
-      // CHECK: [[C_ZERO1:%.*]] = arith.constant {ttg.partition = 2 : i32} 0
-      // CHECK: [[PUT_BUF1:%.*]], [[TOKEN1:%.*]] = nvws.aref.put.enter [[AREF1]][[[C_ZERO1]], [[C_ZERO1]]] {loop.cluster = 1 : i32, loop.stage = 0 : i32, ttg.partition = 2 : i32}
+      // CHECK: [[PUT_BUF1:%.*]], [[TOKEN1:%.*]] = nvws.aref.put.enter [[AREF1]] {loop.cluster = 1 : i32, loop.stage = 0 : i32, ttg.partition = 2 : i32}
       // CHECK-NEXT: nvws.descriptor_load {{.*}} 16384 [[PUT_BUF1]]
-      // CHECK: [[C_ZERO2:%.*]] = arith.constant {ttg.partition = 2 : i32} 0
-      // CHECK: nvws.aref.put.exit [[AREF1]][[[C_ZERO2]]], [[TOKEN1]] [#nvws.async_op<tma_load>] {loop.cluster = 1 : i32, loop.stage = 0 : i32, ttg.partition = 2 : i32}
-      // CHECK: [[C_ZERO3:%.*]] = arith.constant {ttg.partition = 1 : i32} 0
-      // CHECK: [[GET_BUF1:%.*]], [[GET_TOKEN1:%.*]] = nvws.aref.get.enter [[AREF1]][[[C_ZERO3]], [[C_ZERO3]]] {loop.cluster = 0 : i32, loop.stage = 1 : i32, ttg.partition = 1 : i32}
+      // CHECK: nvws.aref.put.exit [[AREF1]], [[TOKEN1]] [#nvws.async_op<tma_load>] {loop.cluster = 1 : i32, loop.stage = 0 : i32, ttg.partition = 2 : i32}
+      // CHECK: [[GET_BUF1:%.*]], [[GET_TOKEN1:%.*]] = nvws.aref.get.enter [[AREF1]] {loop.cluster = 0 : i32, loop.stage = 1 : i32, ttg.partition = 1 : i32}
       %3 = tt.descriptor_load %arg3[%arg1, %2] {loop.cluster = 1 : i32, loop.stage = 0 : i32, ttg.partition = 2 : i32} : !tt.tensordesc<tensor<128x64xf16, #shared>> -> tensor<128x64xf16, #blocked1>
       // CHECK: [[PUT_BUF2:%.*]], [[TOKEN2:%.*]] = nvws.aref.put.enter [[AREF2]]
       // CHECK-NEXT: nvws.descriptor_load {{.*}} 16384 [[PUT_BUF2]]
@@ -51,9 +48,8 @@ module attributes {"ttg.num-warps" = 4 : i32, ttg.target = "cuda:100"} {
       %7 = ttg.memdesc_trans %6 {loop.cluster = 0 : i32, loop.stage = 1 : i32, order = array<i32: 1, 0>, ttg.partition = 1 : i32} : !ttg.memdesc<128x64xf16, #shared, #smem> -> !ttg.memdesc<64x128xf16, #shared1, #smem>
       // CHECK: ttng.tc_gen5_mma [[GET_BUF1]], [[RHS]], {{.*}}, {{.*}}, {{.*}}
       %8 = ttng.tc_gen5_mma %5, %7, %result[%arg6], %true, %true {loop.cluster = 0 : i32, loop.stage = 1 : i32, ttg.partition = 1 : i32} : !ttg.memdesc<128x64xf16, #shared, #smem>, !ttg.memdesc<64x128xf16, #shared1, #smem>, !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>
-      // CHECK: nvws.aref.get.exit [[AREF2]][{{.*}}], [[GET_TOKEN2]]
-      // CHECK: [[C_ZERO4:%.*]] = arith.constant {ttg.partition = 1 : i32} 0
-      // CHECK: nvws.aref.get.exit [[AREF1]][[[C_ZERO4]]], [[GET_TOKEN1]] [#nvws.async_op<tc5mma>] {loop.cluster = 0 : i32, loop.stage = 1 : i32, ttg.partition = 1 : i32}
+      // CHECK: nvws.aref.get.exit [[AREF2]], [[GET_TOKEN2]]
+      // CHECK: nvws.aref.get.exit [[AREF1]], [[GET_TOKEN1]] [#nvws.async_op<tc5mma>] {loop.cluster = 0 : i32, loop.stage = 1 : i32, ttg.partition = 1 : i32}
       scf.yield %8 : !ttg.async.token
     } {tt.num_stages = 2 : i32, tt.scheduled_max_stage = 1 : i32, tt.warp_specialize, ttg.partition.stages = [0 : i32, 1 : i32, 0 : i32], ttg.warp_specialize.tag = 0 : i32}
     %result_0, %token_1 = ttng.tmem_load %result[%1] : !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable> -> tensor<128x128xf32, #blocked>