Repository: intel/intel-xpu-backend-for-triton
Pull request #4152: [tensor-descriptor]: Extend support when tensor descriptor created in control flow
Author: login: etiotto
Created: 2025-05-09T18:07:34Z | Merged: 2025-05-28T12:55:52Z

Enhance  tensor descriptor lowering to support cases where descriptors or pointers are created within control flow constructs.


Diff:
diff --git a/test/Triton/Intel/TensorDescToBlockPointer/basic.mlir b/test/Triton/Intel/TensorDescToBlockPointer/basic.mlir
index a616c39668..e0e7e15eec 100644
--- a/test/Triton/Intel/TensorDescToBlockPointer/basic.mlir
+++ b/test/Triton/Intel/TensorDescToBlockPointer/basic.mlir
@@ -6,20 +6,22 @@ module {
     %c64_i32 = arith.constant 64 : i32
     %c8_i32 = arith.constant 8 : i32
     %0 = arith.extsi %arg2 : i32 to i64
-    %desc = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%0, %c1_i64] : <f32>, <tensor<16x128xf32>>
-    %load = tt.descriptor_load %desc[%c8_i32, %c64_i32] : !tt.tensordesc<tensor<16x128xf32>> -> tensor<16x128xf32>
+    %desc1 = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%0, %c1_i64] : <f32>, <tensor<16x128xf32>>
+    %load1 = tt.descriptor_load %desc1[%c8_i32, %c64_i32] : !tt.tensordesc<tensor<16x128xf32>> -> tensor<16x128xf32>
     tt.return
   }
   // CHECK:      tt.func public @test_load([[PARAM_0:%.+]]: !tt.ptr<f32>, [[PARAM_1:%.+]]: i32, [[PARAM_2:%.+]]: i32) {
   // CHECK-NOT:    tt.make_tensor_descriptor
   // CHECK-NOT:    tt.descriptor_load
+  // CHECK-DAG:    [[CST_0_i32:%.+]] = arith.constant 0 : i32
   // CHECK-DAG:    [[CST_1_i64:%.+]] = arith.constant 1 : i64
   // CHECK-DAG:    [[CST_64_i32:%.+]] = arith.constant 64 : i32
   // CHECK-DAG:    [[CST_8_i32:%.+]] = arith.constant 8 : i32
   // CHECK-DAG:    [[EXTSI_PARAM_1:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
   // CHECK-DAG:    [[EXTSI_PARAM_2:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
-  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2]]], {{\[}}[[EXTSI_PARAM_2]], [[CST_1_i64]]], {{\[}}[[CST_8_i32]], [[CST_64_i32]]] {{.*}} : <tensor<16x128xf32>>
-  // CHECK:        [[LOAD:%.+]] = tt.load [[TENSOR_PTR]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x128xf32>>
+  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2]]], {{\[}}[[EXTSI_PARAM_2]], [[CST_1_i64]]], {{\[}}[[CST_0_i32]], [[CST_0_i32]]] {{.*}} : <tensor<16x128xf32>>
+  // CHECK:        [[TENSOR_PTR1:%.+]] = tt.advance [[TENSOR_PTR]], {{\[}}[[CST_8_i32]], [[CST_64_i32]]] : <tensor<16x128xf32>>
+  // CHECK:        [[LOAD1:%.+]] = tt.load [[TENSOR_PTR1]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x128xf32>>
   // CHECK:        tt.return
   // CHECK:      }
 
@@ -29,21 +31,23 @@ module {
     %c8_i32 = arith.constant 8 : i32
     %cst = arith.constant dense<1.000000e+00> : tensor<16x128xf32>
     %0 = arith.extsi %arg2 : i32 to i64
-    %desc = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%0, %c1_i64] : <f32>, <tensor<16x128xf32>>
-    tt.descriptor_store %desc[%c8_i32, %c64_i32], %cst : !tt.tensordesc<tensor<16x128xf32>>, tensor<16x128xf32>
+    %desc1 = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%0, %c1_i64] : <f32>, <tensor<16x128xf32>>
+    tt.descriptor_store %desc1[%c8_i32, %c64_i32], %cst : !tt.tensordesc<tensor<16x128xf32>>, tensor<16x128xf32>
     tt.return
   }
   // CHECK:      tt.func public @test_store([[PARAM_0:%.+]]: !tt.ptr<f32>, [[PARAM_1:%.+]]: i32, [[PARAM_2:%.+]]: i32) {
   // CHECK-NOT:    tt.make_tensor_descriptor
   // CHECK-NOT:    tt.descriptor_store
+  // CHECK-DAG:    [[CST_0_i32:%.+]] = arith.constant 0 : i32
   // CHECK-DAG:    [[CST_1_i64:%.+]] = arith.constant 1 : i64
   // CHECK-DAG:    [[CST_64_i32:%.+]] = arith.constant 64 : i32
   // CHECK-DAG:    [[CST_8_i32:%.+]] = arith.constant 8 : i32
   // CHECK-DAG:    [[CST:%.+]] = arith.constant dense<1.000000e+00> : tensor<16x128xf32>
   // CHECK-DAG:    [[EXTSI_PARAM_1:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
   // CHECK-DAG:    [[EXTSI_PARAM_2:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
-  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2]]], {{\[}}[[EXTSI_PARAM_2]], [[CST_1_i64]]], {{\[}}[[CST_8_i32]], [[CST_64_i32]]] {{.*}} : <tensor<16x128xf32>>
-  // CHECK:        tt.store [[TENSOR_PTR]], [[CST]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x128xf32>>
+  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2]]], {{\[}}[[EXTSI_PARAM_2]], [[CST_1_i64]]], {{\[}}[[CST_0_i32]], [[CST_0_i32]]] {{.*}} : <tensor<16x128xf32>>
+  // CHECK:        [[TENSOR_PTR1:%.+]] = tt.advance [[TENSOR_PTR]], {{\[}}[[CST_8_i32]], [[CST_64_i32]]] : <tensor<16x128xf32>>
+  // CHECK:        tt.store [[TENSOR_PTR1]], [[CST]]  {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x128xf32>>
   // CHECK:        tt.return
   // CHECK:      }
 }
diff --git a/test/Triton/Intel/TensorDescToBlockPointer/loop.mlir b/test/Triton/Intel/TensorDescToBlockPointer/loop.mlir
index 6ee0f888d1..4bde812ecc 100644
--- a/test/Triton/Intel/TensorDescToBlockPointer/loop.mlir
+++ b/test/Triton/Intel/TensorDescToBlockPointer/loop.mlir
@@ -25,15 +25,16 @@ module {
   // CHECK-DAG:    [[CST_1_i64:%.+]] = arith.constant 1 : i64
   // CHECK-DAG:    [[CST_8_i32:%.+]] = arith.constant 8 : i32
   // CHECK-DAG:    [[CST:%.+]] = arith.constant dense<0.000000e+00> : tensor<16x32xf16>
-  // CHECK-DAG:    [[EXTSI_PARAM_2a:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
-  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} iter_args([[VAR_arg1:%.+]] = {{.*}}, [[VAR_arg2:%.+]] = [[CST]]) -> (!tt.tensordesc<tensor<16x32xf16>>, tensor<16x32xf16>) {
-  // CHECK-DAG:      [[IDX_CAST_1:%.+]] = arith.index_cast [[IV]] : index to i32
-  // CHECK-DAG:      [[EXTSI_PARAM_1:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
-  // CHECK-DAG:      [[EXTSI_PARAM_2b:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
-  // CHECK:          [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2b]]], {{\[}}[[EXTSI_PARAM_2a]], [[CST_1_i64]]], {{\[}}[[CST_8_i32]], [[IDX_CAST_1]]] {{.*}} : <tensor<16x32xf16>>
-  // CHECK:          [[LOAD:%.+]] = tt.load [[TENSOR_PTR]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x32xf16>>
+  // CHECK-DAG:    [[EXTSI_PARAM_1:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
+  // CHECK-DAG:    [[EXTSI_PARAM_2:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
+  // CHECK-DAG:    [[CST_0_i32:%.+]] = arith.constant 0 : i32
+  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2]]], {{\[}}[[EXTSI_PARAM_2]], [[CST_1_i64]]], {{\[}}[[CST_0_i32]], [[CST_0_i32]]] {{.*}} : <tensor<16x32xf16>>
+  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} iter_args([[VAR_arg1:%.+]] = [[TENSOR_PTR]], [[VAR_arg2:%.+]] = [[CST]]) -> (!tt.ptr<tensor<16x32xf16>>, tensor<16x32xf16>) {
+  // CHECK:          [[IDX_CAST:%.+]] = arith.index_cast [[IV]] : index to i32
+  // CHECK:          [[TENSOR_PTR_1:%.+]] = tt.advance [[VAR_arg1]], {{\[}}[[CST_8_i32]], [[IDX_CAST]]] : <tensor<16x32xf16>>
+  // CHECK:          [[LOAD:%.+]] = tt.load [[TENSOR_PTR_1]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x32xf16>>
   // CHECK:          [[ADD:%.+]] = arith.addf [[VAR_arg2]], [[LOAD]] : tensor<16x32xf16>
-  // CHECK:          scf.yield {{.*}}, [[ADD]] : !tt.tensordesc<tensor<16x32xf16>>, tensor<16x32xf16>
+  // CHECK:          scf.yield [[VAR_arg1]], [[ADD]] : !tt.ptr<tensor<16x32xf16>>, tensor<16x32xf16>
   // CHECK:        }
   // CHECK:        tt.return
   // CHECK:      }
@@ -60,12 +61,27 @@ module {
     tt.return
   }
   // CHECK:      tt.func public @load_in_loop2({{.*}}) {
-  // CHECK-NOT:    tt.make_tensor_ptr
-  // CHECK-NOT:    tt.load
-  // CHECK:        tt.make_tensor_descriptor
-  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} -> (!tt.tensordesc<tensor<16x32xf16>>, tensor<16x32xf16>) {
-  // CHECK:          tt.descriptor_load
-  // CHECK:          tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.descriptor_load
+  // CHECK-DAG:    [[CST_1_i64:%.+]] = arith.constant 1 : i64
+  // CHECK-DAG:    [[CST_8_i32:%.+]] = arith.constant 8 : i32
+  // CHECK-DAG:    [[CST:%.+]] = arith.constant dense<0.000000e+00> : tensor<16x32xf16>
+  // CHECK-DAG:    [[EXTSI_PARAM_1:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
+  // CHECK-DAG:    [[EXTSI_PARAM_2:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
+  // CHECK-DAG:    [[CST_0_i32:%.+]] = arith.constant 0 : i32
+  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2]]], {{\[}}[[EXTSI_PARAM_2]], [[CST_1_i64]]], {{\[}}[[CST_0_i32]], [[CST_0_i32]]] {{.*}} : <tensor<16x32xf16>>
+  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} iter_args([[VAR_arg1:%.+]] = [[TENSOR_PTR]], [[VAR_arg2:%.+]] = [[CST]]) -> (!tt.ptr<tensor<16x32xf16>>, tensor<16x32xf16>) {
+  // CHECK:          [[IDX_CAST:%.+]] = arith.index_cast [[IV]] : index to i32
+  // CHECK:          [[TENSOR_PTR_1:%.+]] = tt.advance [[VAR_arg1]], {{\[}}[[CST_8_i32]], [[IDX_CAST]]] : <tensor<16x32xf16>>
+  // CHECK:          [[LOAD:%.+]] = tt.load [[TENSOR_PTR_1]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x32xf16>>
+  // CHECK:          [[ADD:%.+]] = arith.addf [[VAR_arg2]], [[LOAD]] : tensor<16x32xf16>
+  // CHECK-DAG:      [[EXTSI_PARAM_1a:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
+  // CHECK-DAG:      [[EXTSI_PARAM_2a:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
+  // CHECK-DAG:      [[CST_0_i32_1:%.+]] = arith.constant 0 : i32
+  // CHECK:          [[TENSOR_PTR2:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_2a]], [[EXTSI_PARAM_1a]]], {{\[}}[[CST_1_i64]], [[EXTSI_PARAM_2]]], {{\[}}[[CST_0_i32_1]], [[CST_0_i32_1]]] {{.*}} : <tensor<16x32xf16>>
+  // CHECK:          [[CMP:%.+]] = arith.cmpi eq, [[IDX_CAST]], [[CST_8_i32]] : i32
+  // CHECK:          [[TENSOR_PTR3:%.+]] = arith.select [[CMP]], [[VAR_arg1]], [[TENSOR_PTR:%.+]] : !tt.ptr<tensor<16x32xf16>>
+  // CHECK:          scf.yield [[TENSOR_PTR3]], [[ADD]] : !tt.ptr<tensor<16x32xf16>>, tensor<16x32xf16>
   // CHECK:        }
   // CHECK:        tt.return
   // CHECK:      }
@@ -87,10 +103,12 @@ module {
     tt.return
   }
   // CHECK:      tt.func public @load_uses_loop_result({{.*}}) {
-  // CHECK-NOT:    tt.make_tensor_ptr
   // CHECK-NOT:    tt.load
-  // CHECK:        tt.make_tensor_descriptor
-  // CHECK:        tt.descriptor_load
+  // CHECK-NOT:    tt.make_tensor_descriptor
+  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr {{.*}} : <tensor<16x32xf16>>
+  // CHECK:        [[FOR_RES:%.+]] = scf.for [[IV:%.+]] = {{.*}} iter_args([[VAR_arg1:%.+]] = [[TENSOR_PTR]]) -> (!tt.ptr<tensor<16x32xf16>>)
+  // CHECK:        [[TENSOR_PTR1:%.+]] = tt.advance [[FOR_RES]], {{.*}} : <tensor<16x32xf16>>
+  // CHECK:        tt.load [[TENSOR_PTR1]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x32xf16>>
   // CHECK:        tt.return
   // CHECK:      }
 
@@ -115,18 +133,19 @@ module {
   // CHECK:      tt.func public @store_in_loop1([[PARAM_0:%.+]]: !tt.ptr<f16>, [[PARAM_1:%.+]]: i32, [[PARAM_2:%.+]]: i32) {
   // CHECK-NOT:    tt.make_tensor_descriptor
   // CHECK-NOT:    tt.descriptor_store
+  // CHECK-DAG:    [[CST_0_i32:%.+]] = arith.constant 0 : i32
   // CHECK-DAG:    [[CST_1_i64:%.+]] = arith.constant 1 : i64
   // CHECK-DAG:    [[CST_8_i32:%.+]] = arith.constant 8 : i32
   // CHECK-DAG:    [[CST:%.+]] = arith.constant dense<0.000000e+00> : tensor<16x32xf16>
-  // CHECK-DAG:    [[EXTSI_PARAM_2a:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
-  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} iter_args([[VAR_arg1:%.+]] = {{.*}}, [[VAR_arg2:%.+]] = [[CST]]) -> (!tt.tensordesc<tensor<16x32xf16>>, tensor<16x32xf16>) {
-  // CHECK-DAG:      [[IDX_CAST_1:%.+]] = arith.index_cast [[IV]] : index to i32
-  // CHECK-DAG:      [[EXTSI_PARAM_1:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
-  // CHECK-DAG:      [[EXTSI_PARAM_2b:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
-  // CHECK:          [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2b]]], {{\[}}[[EXTSI_PARAM_2a]], [[CST_1_i64]]], {{\[}}[[CST_8_i32]], [[IDX_CAST_1]]] {{.*}} : <tensor<16x32xf16>>
-  // CHECK:          tt.store [[TENSOR_PTR]], [[VAR_arg2]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x32xf16>>
+  // CHECK-DAG:    [[EXTSI_PARAM_1:%.+]] = arith.extsi [[PARAM_1]] : i32 to i64
+  // CHECK-DAG:    [[EXTSI_PARAM_2:%.+]] = arith.extsi [[PARAM_2]] : i32 to i64
+  // CHECK:        [[TENSOR_PTR:%.+]] = tt.make_tensor_ptr [[PARAM_0]], {{\[}}[[EXTSI_PARAM_1]], [[EXTSI_PARAM_2]]], {{\[}}[[EXTSI_PARAM_2]], [[CST_1_i64]]], {{\[}}[[CST_0_i32]], [[CST_0_i32]]] {{.*}} : <tensor<16x32xf16>>
+  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} iter_args([[VAR_arg1:%.+]] = [[TENSOR_PTR]], [[VAR_arg2:%.+]] = [[CST]]) -> (!tt.ptr<tensor<16x32xf16>>, tensor<16x32xf16>) {
+  // CHECK:          [[IDX_CAST_1:%.+]] = arith.index_cast [[IV]] : index to i32
+  // CHECK:          [[TENSOR_PTR_1:%.+]] = tt.advance [[VAR_arg1]], {{\[}}[[CST_8_i32]], [[IDX_CAST]]] : <tensor<16x32xf16>>
+  // CHECK:          tt.store [[TENSOR_PTR_1]], [[VAR_arg2]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<16x32xf16>>
   // CHECK:          [[ADD:%.+]] = arith.addf [[VAR_arg2]], [[CST]] : tensor<16x32xf16>
-  // CHECK:          scf.yield {{.*}}, [[ADD]] : !tt.tensordesc<tensor<16x32xf16>>, tensor<16x32xf16>
+  // CHECK:          scf.yield [[VAR_arg1]], [[ADD]] : !tt.ptr<tensor<16x32xf16>>, tensor<16x32xf16>
   // CHECK:        }
   // CHECK:        tt.return
   // CHECK:      }
@@ -153,12 +172,12 @@ module {
     tt.return
   }
   // CHECK:      tt.func public @store_in_loop2({{.*}}) {
-  // CHECK-NOT:    tt.make_tensor_ptr
-  // CHECK-NOT:    tt.store
-  // CHECK:        tt.make_tensor_descriptor
-  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} -> (!tt.tensordesc<tensor<16x32xf16>>, tensor<16x32xf16>) {
-  // CHECK:          tt.descriptor_store
-  // CHECK:          tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.descriptor_store
+  // CHECK:        tt.make_tensor_ptr
+  // CHECK:        [[FOR_RES:%.+]]:2 = scf.for [[IV:%.+]] = {{.*}} -> (!tt.ptr<tensor<16x32xf16>>, tensor<16x32xf16>) {
+  // CHECK:          tt.advance
+  // CHECK:          tt.store
   // CHECK:        }
   // CHECK:        tt.return
   // CHECK:      }
@@ -181,11 +200,123 @@ module {
     tt.return
   }
   // CHECK:      tt.func public @store_uses_loop_result({{.*}}) {
-  // CHECK-NOT:    tt.make_tensor_ptr
-  // CHECK-NOT:    tt.store
-  // CHECK:        tt.make_tensor_descriptor
-  // CHECK:        tt.descriptor_store
+  // CHECK-NOT:    tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.descriptor_store
+  // CHECK:        tt.make_tensor_ptr
+  // CHECK:        tt.advance
+  // CHECK:        tt.store
   // CHECK:        tt.return
   // CHECK:      }
 
+  // COM: While loop contains a descriptor load operation.
+  tt.func public @load_in_while_loop(%arg0: !tt.ptr<f32>, %arg1: i32, %arg2: i32) {
+    %c0_i32 = arith.constant 0 : i32
+    %c1_i64 = arith.constant 1 : i64
+    %0 = tt.get_program_id x : i32
+    %3 = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%c1_i64, %c1_i64] : <f32>, <tensor<8x128xf32>>
+    %5 = scf.while (%arg3 = %3) : (!tt.tensordesc<tensor<8x128xf32>>) -> (!tt.tensordesc<tensor<8x128xf32>>) {
+      %6 = arith.cmpi slt, %c0_i32, %arg2 : i32
+      scf.condition(%6) %arg3 : !tt.tensordesc<tensor<8x128xf32>>
+    } do {
+    ^bb0(%arg3: !tt.tensordesc<tensor<8x128xf32>>):
+      %12 = tt.descriptor_load %arg3[%0, %c0_i32] : !tt.tensordesc<tensor<8x128xf32>> -> tensor<8x128xf32>
+      scf.yield %arg3 : !tt.tensordesc<tensor<8x128xf32>>
+    }
+    tt.return
+  }
+  // CHECK: tt.func public @load_in_while_loop({{.*}}) {
+  // CHECK-NOT:    tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.descriptor_load
+  // CHECK:        [[TENSOR_PTR:%.*]] = tt.make_tensor_ptr {{.*}} : <tensor<8x128xf32>
+  // CHECK:        scf.while ([[ARG3:%.*]] = [[TENSOR_PTR]]) : (!tt.ptr<tensor<8x128xf32>>) -> !tt.ptr<tensor<8x128xf32>> {
+  // CHECK:          scf.condition({{.*}}) [[ARG3]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:        } do {
+  // CHECK:        ^bb0([[ARG4:%.*]]: !tt.ptr<tensor<8x128xf32>>):
+  // CHECK:          [[PTR1:%.*]] = tt.advance [[ARG4]], {{.*}} : <tensor<8x128xf32>
+  // CHECK:          tt.load [[PTR1]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:          scf.yield [[ARG4]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:        }
+
+  // COM: For loop yields a tensor descriptor used by a while loop.
+  tt.func public @while_uses_tdesc_yielded_by_for_loop(%arg0: !tt.ptr<f32>, %arg1: i32, %arg2: i32) {
+    %c1_i64 = arith.constant 1 : i64
+    %c128_i32 = arith.constant 128 : i32
+    %c0_i32 = arith.constant 0 : i32
+    %c8_i32 = arith.constant 8 : i32
+    %0 = tt.get_program_id x : i32
+    %2 = arith.extsi %arg2 : i32 to i64
+    %3 = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%2, %c1_i64] : <f32>, <tensor<8x128xf32>>
+    %4 = scf.for %arg3 = %c0_i32 to %arg2 step %c128_i32 iter_args(%arg4 = %3) -> (!tt.tensordesc<tensor<8x128xf32>>) : i32 {
+      scf.yield %arg4 : !tt.tensordesc<tensor<8x128xf32>>
+    }
+    %5 = scf.while (%arg3 = %4) : (!tt.tensordesc<tensor<8x128xf32>>) -> (!tt.tensordesc<tensor<8x128xf32>>) {
+      %6 = arith.cmpi slt, %0, %arg2 : i32
+      scf.condition(%6) %arg3 : !tt.tensordesc<tensor<8x128xf32>>
+    } do {
+    ^bb0(%arg3: !tt.tensordesc<tensor<8x128xf32>>):
+      %12 = tt.descriptor_load %arg3[%c8_i32, %c8_i32] : !tt.tensordesc<tensor<8x128xf32>> -> tensor<8x128xf32>
+      scf.yield %arg3 : !tt.tensordesc<tensor<8x128xf32>>
+    }
+    tt.return
+  }
+  // CHECK:      tt.func public @while_uses_tdesc_yielded_by_for_loop({{.*}}) {
+  // CHECK-NOT:    tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.descriptor_load
+  // CHECK:        [[TENSOR_PTR:%.*]] = tt.make_tensor_ptr {{.*}} : <tensor<8x128xf32>
+  // CHECK:        [[FOR_RES:%.+]] = scf.for [[IV:%.+]] = {{.*}} iter_args([[ARG3:%.*]] = [[TENSOR_PTR]]) -> (!tt.ptr<tensor<8x128xf32>>) : i32 {
+  // CHECK:          scf.yield {{.*}} : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:        }
+  // CHECK:        scf.while ([[ARG3:%.*]] = [[FOR_RES]]) : (!tt.ptr<tensor<8x128xf32>>) -> !tt.ptr<tensor<8x128xf32>> {
+  // CHECK:          scf.condition({{.*}}) [[ARG3]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:        } do {
+  // CHECK:        ^bb0([[ARG4:%.*]]: !tt.ptr<tensor<8x128xf32>>):
+  // CHECK:          [[TENSOR_PTR1:%.*]] = tt.advance [[ARG4]], {{.*}} : <tensor<8x128xf32>
+  // CHECK:          tt.load [[TENSOR_PTR1]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:          scf.yield [[ARG4]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:        }
+
+  // COM: While loop containing a if statement yielding a tensor descriptor used by a load operation.
+  tt.func public @while_loop_with_if_stmt(%arg0: !tt.ptr<f32>, %arg1: i32, %arg2: i32) {
+    %c1_i64 = arith.constant 1 : i64
+    %c128_i32 = arith.constant 128 : i32
+    %c0_i32 = arith.constant 0 : i32
+    %0 = tt.get_program_id x : i32
+    %3 = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%c1_i64, %c1_i64] : <f32>, <tensor<8x128xf32>>
+    %5 = scf.while (%arg3 = %3) : (!tt.tensordesc<tensor<8x128xf32>>) -> (!tt.tensordesc<tensor<8x128xf32>>) {
+      %6 = arith.cmpi slt, %0, %c128_i32 : i32
+      scf.condition(%6) %arg3 : !tt.tensordesc<tensor<8x128xf32>>
+    } do {
+    ^bb0(%arg3: !tt.tensordesc<tensor<8x128xf32>>):
+      %7 = arith.cmpi eq, %0, %c0_i32 : i32
+      %10 = arith.select %7, %c1_i64, %c1_i64 : i64
+      %11 = scf.if %7 -> (!tt.tensordesc<tensor<8x128xf32>>) {
+        %16 = tt.make_tensor_descriptor %arg0, [%arg1, %arg2], [%c1_i64, %c1_i64] : <f32>, <tensor<8x128xf32>>
+        scf.yield %16 : !tt.tensordesc<tensor<8x128xf32>>
+      } else {
+        scf.yield %arg3 : !tt.tensordesc<tensor<8x128xf32>>
+      }
+      %12 = tt.descriptor_load %11[%c128_i32, %c128_i32] : !tt.tensordesc<tensor<8x128xf32>> -> tensor<8x128xf32>
+      scf.yield %11 : !tt.tensordesc<tensor<8x128xf32>>
+    }
+    tt.return
+  }
+  // CHECK:      tt.func public @while_loop_with_if_stmt({{.*}}) {
+  // CHECK-NOT:    tt.make_tensor_descriptor
+  // CHECK-NOT:    tt.descriptor_load
+  // CHECK:        [[TENSOR_PTR:%.*]] = tt.make_tensor_ptr {{.*}} : <tensor<8x128xf32>
+  // CHECK:        scf.while ([[ARG3:%.*]] = [[TENSOR_PTR]]) : (!tt.ptr<tensor<8x128xf32>>) -> !tt.ptr<tensor<8x128xf32>> {
+  // CHECK:          scf.condition({{.*}}) [[ARG3]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:        } do {
+  // CHECK:        ^bb0([[ARG4:%.*]]: !tt.ptr<tensor<8x128xf32>>):
+  // CHECK:          [[IF_RES:%.*]] = scf.if {{.*}} -> (!tt.ptr<tensor<8x128xf32>>) {
+  // CHECK:            [[TENSOR_PTR1:%.*]] = tt.make_tensor_ptr {{.*}} : <tensor<8x128xf32>
+  // CHECK:            scf.yield [[TENSOR_PTR1]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:          } else {
+  // CHECK:            scf.yield [[ARG4]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:          }
+  // CHECK:          [[TENSOR_PTR2:%.*]] = tt.advance [[IF_RES]], {{.*}} : <tensor<8x128xf32>
+  // CHECK:          tt.load [[TENSOR_PTR2]] {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:          scf.yield [[IF_RES]] : !tt.ptr<tensor<8x128xf32>>
+  // CHECK:        }
+
 }
diff --git a/test/TritonIntelGPU/backward_combine_dpas_dot_layout.mlir b/test/TritonIntelGPU/backward_combine_dpas_dot_layout.mlir
index 5ca2366b30..91487e4bb3 100644
--- a/test/TritonIntelGPU/backward_combine_dpas_dot_layout.mlir
+++ b/test/TritonIntelGPU/backward_combine_dpas_dot_layout.mlir
@@ -66,7 +66,6 @@ module attributes {"ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 16 : i32,
     }
     %24 = arith.truncf %23#0 : tensor<64x256xf32, #dpas> to tensor<64x256xf16, #dpas>
     %25 = ttg.convert_layout %24 : tensor<64x256xf16, #dpas> -> tensor<64x256xf16, #blocked1>
-
     %27 = tt.make_tensor_ptr %arg2, [%15, %20], [%arg8, %c1_i64], [%14, %19] {order = array<i32: 1, 0>} : <tensor<64x256xf16, #blocked1>>
     tt.store %27, %25 {boundaryCheck = array<i32: 0, 1>} : !tt.ptr<tensor<64x256xf16, #blocked1>>
     tt.return
diff --git a/third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp b/third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
index c4f238c2c0..5efc369b61 100644
--- a/third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
+++ b/third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
@@ -1,13 +1,14 @@
 #include "intel/include/Dialect/Triton/Transforms/Passes.h"
+#include "intel/include/Dialect/TritonGEN/IR/TritonGENDialect.h"
+#include "intel/include/Utils/Utility.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/Types.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
-#include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "triton-intel-tdesc-to-block-pointer"
@@ -62,6 +63,12 @@ struct TritonIntelTensorDescToBlockPointer
 
     moduleOp->walk<WalkOrder::PreOrder>([&](Operation *op) {
       return TypeSwitch<Operation *, WalkResult>(op)
+          .Case<tt::MakeTensorDescOp>([&](auto makeTensorDescOp) {
+            if (failed(rewriteMakeTensorDescriptorOp(makeTensorDescOp)))
+              makeTensorDescOp->emitRemark(
+                  "TritonIntelTensorDescToBlockPointer: Failed to rewrite");
+            return WalkResult::advance();
+          })
           .Case<tt::DescriptorLoadOp, tt::DescriptorStoreOp>(
               [&](auto loadOrStoreOp) {
                 if (failed(rewriteDescriptorLoadOrStoreOp(loadOrStoreOp)))
@@ -77,56 +84,16 @@ struct TritonIntelTensorDescToBlockPointer
   }
 
 private:
-  tt::MakeTensorDescOp getMakeTensorDescOp(Value base) const {
-    assert(base && isa<tt::TensorDescType>(base.getType()) &&
-           "Expecting tensor desc");
-
-    Operation *defOp = base.getDefiningOp();
-    if (!defOp) {
-      BlockArgument blockArg = cast<BlockArgument>(base);
-      Operation *parentOp = blockArg.getOwner()->getParentOp();
-      if (scf::ForOp forOp = dyn_cast<scf::ForOp>(parentOp)) {
-        unsigned numIVs = forOp.getNumInductionVars();
-        int initArgIdx = blockArg.getArgNumber() - numIVs;
-        if (isModifiedInLoop(forOp, blockArg)) {
-          LLVM_DEBUG(llvm::dbgs() << blockArg << "is loop variant");
-          return nullptr;
-        }
-        Operation::operand_range initArgs = forOp.getInitArgs();
-        assert(initArgIdx >= 0 && initArgIdx < initArgs.size() &&
-               "Unexpected 'initArgIdx' value");
-        return getMakeTensorDescOp(initArgs[initArgIdx]);
-      }
-      LLVM_DEBUG(llvm::dbgs()
-                 << "TODO: Unhandled non operation: " << base << "\n");
-      return nullptr;
-    }
-
-    if (defOp->getNumRegions() != 0) {
-      LLVM_DEBUG(llvm::dbgs() << "TODO: defOp with region: " << *defOp << "\n");
-      return nullptr;
-    }
-    if (auto makeTensorDescOp = dyn_cast<tt::MakeTensorDescOp>(defOp))
-      return makeTensorDescOp;
-
-    llvm_unreachable("TODO: Unhandled defOp kind");
-    return nullptr;
-  }
-
-  bool isModifiedInLoop(scf::ForOp forOp, BlockArgument &blockArg) const {
-    unsigned argNo = blockArg.getArgNumber();
-    unsigned numIVs = forOp.getNumInductionVars();
-    int initArgIdx = blockArg.getArgNumber() - numIVs;
-    Value yieldedVal = forOp.getYieldedValues()[initArgIdx];
-    return (yieldedVal != blockArg);
-  }
-
-  Value findOrCreateMakeTensorPtr(Location loc, Value base, ValueRange shape,
-                                  ValueRange strides, ValueRange offsets,
-                                  ArrayRef<int32_t> sizes, OpBuilder &builder) {
+  // Create a new block pointer if a suitable one doesn't already exist.
+  // Otherwise, return the existing one. The function takes the base, shape,
+  // strides, offsets, sizes of the block pointer to create/lookup and its
+  // tensor element type (to ensure the block pointer has the tensor layout).
+  tt::MakeTensorPtrOp
+  findOrCreateMakeTensorPtr(Location loc, Value base, ValueRange shape,
+                            ValueRange strides, ValueRange offsets,
+                            ArrayRef<int32_t> sizes, OpBuilder &builder) {
     Block *block = builder.getInsertionBlock();
     const Block::iterator insertPoint = builder.getInsertionPoint();
-
     auto it = std::find_if(block->begin(), insertPoint, [&](Operation &op) {
       if (auto makeTensorPtrOp = dyn_cast<tt::MakeTensorPtrOp>(op)) {
         triton::PointerType resType = makeTensorPtrOp.getResult().getType();
@@ -138,6 +105,7 @@ struct TritonIntelTensorDescToBlockPointer
           }
           return true;
         };
+
         return makeTensorPtrOp.getBase() == base &&
                makeTensorPtrOp.getShape() == shape &&
                makeTensorPtrOp.getStrides() == strides &&
@@ -147,81 +115,136 @@ struct TritonIntelTensorDescToBlockPointer
       return false;
     });
 
+    auto makeTensorPtrOp = [&]() {
+      auto makeTensorPtr = builder.create<tt::MakeTensorPtrOp>(
+          loc, base, shape, strides, offsets, sizes,
+          builder.getDenseI32ArrayAttr({1, 0}));
+      return makeTensorPtr;
+    };
+
     return (it != insertPoint) ? cast<tt::MakeTensorPtrOp>(*it)
-                               : builder.createOrFold<tt::MakeTensorPtrOp>(
-                                     loc, base, shape, strides, offsets, sizes,
-                                     builder.getDenseI32ArrayAttr({1, 0}));
+                               : makeTensorPtrOp();
   }
 
-  template <typename OpTy,
-            std::enable_if_t<llvm::is_one_of<OpTy, tt::DescriptorLoadOp,
-                                             tt::DescriptorStoreOp>::value,
-                             bool> = true>
-  LogicalResult rewriteDescriptorLoadOrStoreOp(OpTy op) {
+  void propagateToLoops(Operation *op) {
+    if (auto loopOp = dyn_cast<LoopLikeOpInterface>(op)) {
+      bool updated = false;
+      for (auto [initArg, rgnInitArg, yieldVal, loopRes] :
+           llvm::zip(loopOp.getInits(), loopOp.getRegionIterArgs(),
+                     loopOp.getYieldedValues(), loopOp->getResults())) {
+        Type initArgType = initArg.getType();
+        Type rgnInitArgType = rgnInitArg.getType();
+        assert(rgnInitArgType == loopRes.getType() &&
+               rgnInitArgType == yieldVal.getType() && "Type mismatch");
+        if (rgnInitArgType != initArgType) {
+          rgnInitArg.setType(initArgType);
+          yieldVal.setType(initArgType);
+          loopRes.setType(initArgType);
+          updated = true;
+        }
+      }
+      if (!updated)
+        return;
+
+      // For while loops we also need to update the "after" region arguments.
+      if (auto loopOp = dyn_cast<scf::WhileOp>(op)) {
+        for (auto [initArg, rgnAfterArg] :
+             llvm::zip(loopOp.getInits(), loopOp.getAfterArguments())) {
+          Type initArgType = initArg.getType();
+          if (rgnAfterArg.getType() != initArgType)
+            rgnAfterArg.setType(initArgType);
+        }
+      }
+
+      // Propagate the loop results to their users.
+      for (Operation *user : loopOp->getUsers())
+        propagateToLoops(user);
+    }
+  }
+
+  LogicalResult rewriteMakeTensorDescriptorOp(tt::MakeTensorDescOp op) {
     assert(op && "Expecting a valid operation");
     LLVM_DEBUG(llvm::dbgs() << "Rewriting: " << op << "\n");
 
     OpBuilder builder(op);
     Location loc = op.getLoc();
-    TypedValue<tt::TensorDescType> tDesc = op.getDesc();
-    tt::TensorDescType tDescType = tDesc.getType();
-    tt::MakeTensorDescOp makeTensorDescOp = getMakeTensorDescOp(tDesc);
-
-    if (!makeTensorDescOp) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "could not find tt.make_tensor_descriptor defining: "
-                 << tDesc << "\n");
-      return failure();
-    }
-
-    LLVM_DEBUG(llvm::dbgs() << "which has tdesc: " << makeTensorDescOp << "\n");
+    tt::TensorDescType tDescType = op.getType();
 
     // Create a new block pointer if a suitable one doesn't already exist.
     SmallVector<Value> shapes, strides, offsets;
     SmallVector<int32_t> sizes;
-    for (const auto [shape, stride, offset, size] :
-         llvm::zip(makeTensorDescOp.getShape(), makeTensorDescOp.getStrides(),
-                   op.getIndices(), tDescType.getBlockType().getShape())) {
+    for (const auto [shape, stride, size] :
+         llvm::zip(op.getShape(), op.getStrides(),
+                   tDescType.getBlockType().getShape())) {
       shapes.push_back(findOrCreateCast(
           loc, shape, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
       strides.push_back(findOrCreateCast(
           loc, stride, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
-      offsets.push_back(findOrCreateCast(
-          loc, offset, builder.getIntegerType(offsetBitwidth), builder));
+      Value zero =
+          tt::intel::findOrCreateIntConstant(loc, 0, offsetBitwidth, builder);
+      offsets.push_back(zero);
       sizes.push_back(static_cast<int32_t>(size));
     }
 
-    Value makeTensorPtrOp =
-        findOrCreateMakeTensorPtr(loc, makeTensorDescOp.getBase(), shapes,
-                                  strides, offsets, sizes, builder);
-
+    auto tensorPtr = findOrCreateMakeTensorPtr(
+        loc, op.getBase(), shapes, strides, offsets, sizes, builder);
     LLVM_DEBUG({
       llvm::dbgs() << "With:\n";
-      llvm::dbgs().indent(2) << makeTensorPtrOp << "\n";
+      llvm::dbgs().indent(2) << tensorPtr << "\n";
     });
 
+    op->replaceAllUsesWith(tensorPtr);
+    cleanUp.insert(op);
+
+    // Propagate the `tensorPtr` type to loops init args, yielded values,
+    // results, ... (if necessary).
+    for (Operation *user : tensorPtr->getUsers())
+      propagateToLoops(user);
+
+    return success();
+  }
+
+  template <typename OpTy,
+            std::enable_if_t<llvm::is_one_of<OpTy, tt::DescriptorLoadOp,
+                                             tt::DescriptorStoreOp>::value,
+                             bool> = true>
+  LogicalResult rewriteDescriptorLoadOrStoreOp(OpTy op) {
+    assert(op && "Expecting a valid operation");
+    LLVM_DEBUG(llvm::dbgs() << "Rewriting: " << op << "\n");
+
+    OpBuilder builder(op);
+    Location loc = op.getLoc();
+    Value ptr = op.getOperand(0);
+    assert(triton::isTensorPointerType(ptr.getType()) &&
+           "Expecting a block ptr");
+    auto ptrType = cast<tt::PointerType>(ptr.getType());
+    auto tensorType = cast<RankedTensorType>(ptrType.getPointeeType());
+
+    ptr =
+        builder.create<tt::AdvanceOp>(loc, ptr.getType(), ptr, op.getIndices());
+
     SmallVector<int32_t> boundaryCheck;
-    for (size_t i = 0; i < makeTensorDescOp.getShape().size(); ++i)
+    for (size_t i = 0; i < tensorType.getRank(); ++i)
       boundaryCheck.push_back(i);
+
     constexpr bool isLoad = std::is_same_v<OpTy, tt::DescriptorLoadOp>;
     if constexpr (isLoad) {
       auto loadOp = builder.createOrFold<tt::LoadOp>(
-          loc, makeTensorPtrOp, boundaryCheck, /*padding*/ std::nullopt,
-          op.getCache(), op.getEvict(),
+          loc, ptr, boundaryCheck,
+          /*padding*/ std::nullopt, op.getCache(), op.getEvict(),
           /*volatile*/ false);
       LLVM_DEBUG(llvm::dbgs().indent(2) << loadOp << "\n");
       op.replaceAllUsesWith(loadOp);
     } else {
       [[maybe_unused]] auto storeOp = builder.createOrFold<tt::StoreOp>(
-          loc, makeTensorPtrOp, op.getSrc(), boundaryCheck,
-          tt::CacheModifier::NONE, tt::EvictionPolicy::NORMAL);
+          loc, ptr, op.getSrc(), boundaryCheck, tt::CacheModifier::NONE,
+          tt::EvictionPolicy::NORMAL);
       LLVM_DEBUG(llvm::dbgs().indent(2) << storeOp << "\n");
     }
 
     cleanUp.insert(op);
-    cleanUp.insert(makeTensorDescOp);
 
     return success();
   }


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR aims to ensure that the block pointer created via tt.make_tensor_ptr correctly inherits the layout from the tensor descriptor load operation. Key changes include adding a TritonGEN header, modifying the findOrCreateMakeTensorPtr function to incorporate the tensor element type, and adding a new lambda for pointer type creation.

<details>
<summary>Files not reviewed (1)</summary>

* **test/Triton/Intel/TensorDescToBlockPointer/basic.mlir**: Language not supported
</details>



state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T18:08:00Z
user: login: chengjunlu
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T23:19:26Z
user: login: mfrancepillois
body: The code LGTM, but I'm not sure to understand why we need this PR if the conversion from `tensor_descriptor` to ` block_pointer` is one of the first passes to run at the `ttir` level, how can we have `tensor_descriptor` ops with different layouts given the encodings are assigned to ops when lowering from the `ttir` level to the `ttgir` level.
As I understand it, we are not supposed to have `tensor_descriptor` ops after translating to `ttgir` dialect, no?
state: APPROVED
author_association: NONE
submitted_at: 2025-05-12T08:59:01Z
user: login: whitneywhtsang
state: DISMISSED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T17:41:05Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

Enhance layout propagation and tensor descriptor lowering to support cases where descriptors or pointers are created within control flow constructs.

- Add `updateAdvanceOpChain` to recursively update chains of `AdvanceOp` users.
- Refactor `rewriteStoreOp` to trace back through `AdvanceOp` chains before creating new `MakeTensorPtrOp`.
- Rewrite the TensorDesc-to-block-pointer pass to drop old descriptor lookup, always find or create `MakeTensorPtrOp`, and handle loop-carried block pointer types.

### Reviewed Changes

Copilot reviewed 2 out of 5 changed files in this pull request and generated 2 comments.

| File                                                                                                                                      | Description                                                      |
| ----------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- |
| third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp                                                                | Added recursive chain update, improved `rewriteStoreOp`, and added verification asserts after each transformation stage. |
| third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp                                                               | Removed legacy descriptor lookup, consolidated pointer creation via `findOrCreateMakeTensorPtr`, replaced offset logic, and updated loop argument types. |

<details>
<summary>Files not reviewed (3)</summary>

* **test/Triton/Intel/TensorDescToBlockPointer/basic.mlir**: Language not supported
* **test/Triton/Intel/TensorDescToBlockPointer/loop.mlir**: Language not supported
* **test/TritonIntelGPU/backward_combine_dpas_dot_layout.mlir**: Language not supported
</details>

<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp:793**
* The variable `value` is undefined in this scope, leading to a compile error. It should reference the store's value (e.g., `storeOp.getValue()`) or the converted value extracted earlier.
```
Value dataToStore = getValueAs(value, encoding);
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T16:29:42Z
user: login: etiotto
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T17:04:18Z
user: login: whitneywhtsang
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T21:36:24Z
user: login: chengjunlu
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T00:50:49Z
user: login: etiotto
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T15:32:57Z
user: login: etiotto
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T15:34:48Z
user: login: whitneywhtsang
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T18:30:31Z
user: login: mfrancepillois
state: APPROVED
author_association: NONE
submitted_at: 2025-05-27T12:57:34Z
user: login: etiotto
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T12:55:00Z

Review comments:
diff_hunk: @@ -73,6 +74,8 @@ struct TritonIntelTensorDescToBlockPointer
     });
 
     finalize();
+
+    llvm::errs() << "moduleOp: " << *moduleOp << "\n";
path: third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
user: login: Copilot
body: The debug output on this line appears to be left in production code. It is recommended to remove or guard this debug statement to avoid unintended logging in release builds.
```suggestion

#ifdef DEBUG
    llvm::errs() << "moduleOp: " << *moduleOp << "\n";
#endif
```
created_at: 2025-05-09T18:08:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -176,6 +195,11 @@ struct TritonIntelTensorDescToBlockPointer
 
     LLVM_DEBUG(llvm::dbgs() << "which has tdesc: " << makeTensorDescOp << "\n");
 
+    auto createPointerType = [](RankedTensorType tensorType) {
+      return tt::PointerType::get(
+          tensorType, tt::TritonGEN::TritonGENMemorySpace::kCrossWorkgroup);
+    };
path: third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
user: login: Copilot
body: The lambda 'createPointerType' is defined but not used anywhere, leading to dead code. It is recommended to remove this unused lambda if it is not intended for future use.
```suggestion
// Removed unused lambda `createPointerType`.
```
created_at: 2025-05-09T18:08:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -692,6 +692,26 @@ void LayoutPropagation::rewriteAssertOp(AssertOp assertOp) {
   assertOp->setOperand(0, newOperand);
 }
 
+// Recursively update the operands in a chain of AdvanceOps, after setting the
+// pointer operand of the first one.
+static void updateAdvanceOpChain(AdvanceOp advanceOp, Value makeTensorPtrOp,
+                                 Value dataToStore) {
+  OpBuilder rewriter(advanceOp);
+  auto newAdvanceOp =
+      rewriter.create<AdvanceOp>(advanceOp.getLoc(), makeTensorPtrOp.getType(),
+                                 makeTensorPtrOp, advanceOp.getOffsets());
+
+  SmallVector<Operation *> advanceOpUsers(advanceOp->getUsers());
+  for (Operation *user : advanceOpUsers) {
+    if (auto storeOp = dyn_cast<StoreOp>(user)) {
+      storeOp.setOperand(0, newAdvanceOp);
+      storeOp.setOperand(1, dataToStore);
+    } else if (auto advanceOp = dyn_cast<AdvanceOp>(user)) {
+      updateAdvanceOpChain(advanceOp, makeTensorPtrOp, dataToStore);
path: third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
user: login: Copilot
body: This recursive call always passes the original `makeTensorPtrOp`, but should use the newly created `newAdvanceOp` so that subsequent stores and advances chain the correct pointer.
```suggestion
      updateAdvanceOpChain(advanceOp, newAdvanceOp, dataToStore);
```
created_at: 2025-05-20T16:29:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -147,77 +113,106 @@ struct TritonIntelTensorDescToBlockPointer
       return false;
     });
 
+    auto makeTensorPtrOp = [&]() {
+      Value makeTensorPtr = builder.create<tt::MakeTensorPtrOp>(
+          loc, base, shape, strides, offsets, sizes,
+          builder.getDenseI32ArrayAttr({1, 0}));
+      return makeTensorPtr;
+    };
+
     return (it != insertPoint) ? cast<tt::MakeTensorPtrOp>(*it)
-                               : builder.createOrFold<tt::MakeTensorPtrOp>(
-                                     loc, base, shape, strides, offsets, sizes,
-                                     builder.getDenseI32ArrayAttr({1, 0}));
+                               : makeTensorPtrOp();
   }
 
-  template <typename OpTy,
-            std::enable_if_t<llvm::is_one_of<OpTy, tt::DescriptorLoadOp,
-                                             tt::DescriptorStoreOp>::value,
-                             bool> = true>
-  LogicalResult rewriteDescriptorLoadOrStoreOp(OpTy op) {
+  LogicalResult rewriteMakeTensorDescriptorOp(tt::MakeTensorDescOp op) {
     assert(op && "Expecting a valid operation");
     LLVM_DEBUG(llvm::dbgs() << "Rewriting: " << op << "\n");
 
     OpBuilder builder(op);
     Location loc = op.getLoc();
-    TypedValue<tt::TensorDescType> tDesc = op.getDesc();
-    tt::TensorDescType tDescType = tDesc.getType();
-    tt::MakeTensorDescOp makeTensorDescOp = getMakeTensorDescOp(tDesc);
-
-    if (!makeTensorDescOp) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "could not find tt.make_tensor_descriptor defining: "
-                 << tDesc << "\n");
-      return failure();
-    }
-
-    LLVM_DEBUG(llvm::dbgs() << "which has tdesc: " << makeTensorDescOp << "\n");
+    tt::TensorDescType tDescType = op.getType();
 
     // Create a new block pointer if a suitable one doesn't already exist.
     SmallVector<Value> shapes, strides, offsets;
     SmallVector<int32_t> sizes;
-    for (const auto [shape, stride, offset, size] :
-         llvm::zip(makeTensorDescOp.getShape(), makeTensorDescOp.getStrides(),
-                   op.getIndices(), tDescType.getBlockType().getShape())) {
+    for (const auto [shape, stride, size] :
+         llvm::zip(op.getShape(), op.getStrides(),
+                   tDescType.getBlockType().getShape())) {
       shapes.push_back(findOrCreateCast(
           loc, shape, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
       strides.push_back(findOrCreateCast(
           loc, stride, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
-      offsets.push_back(findOrCreateCast(
-          loc, offset, builder.getIntegerType(offsetBitwidth), builder));
+      Value zero =
+          tt::intel::findOrCreateIntConstant(loc, 0, offsetBitwidth, builder);
+      offsets.push_back(zero);
path: third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
user: login: Copilot
body: Always using zero for offsets drops the original descriptor indices; this loses intended offsets. You should preserve or compute offsets from `op.getIndices()` rather than defaulting to zero.
```suggestion
      Value index = findOrCreateCast(
          loc, op.getIndices()[offsets.size()], builder.getIntegerType(offsetBitwidth),
          builder);
      offsets.push_back(index);
```
created_at: 2025-05-20T16:29:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -147,77 +113,106 @@ struct TritonIntelTensorDescToBlockPointer
       return false;
     });
 
+    auto makeTensorPtrOp = [&]() {
+      Value makeTensorPtr = builder.create<tt::MakeTensorPtrOp>(
+          loc, base, shape, strides, offsets, sizes,
+          builder.getDenseI32ArrayAttr({1, 0}));
+      return makeTensorPtr;
+    };
+
     return (it != insertPoint) ? cast<tt::MakeTensorPtrOp>(*it)
-                               : builder.createOrFold<tt::MakeTensorPtrOp>(
-                                     loc, base, shape, strides, offsets, sizes,
-                                     builder.getDenseI32ArrayAttr({1, 0}));
+                               : makeTensorPtrOp();
   }
 
-  template <typename OpTy,
-            std::enable_if_t<llvm::is_one_of<OpTy, tt::DescriptorLoadOp,
-                                             tt::DescriptorStoreOp>::value,
-                             bool> = true>
-  LogicalResult rewriteDescriptorLoadOrStoreOp(OpTy op) {
+  LogicalResult rewriteMakeTensorDescriptorOp(tt::MakeTensorDescOp op) {
     assert(op && "Expecting a valid operation");
     LLVM_DEBUG(llvm::dbgs() << "Rewriting: " << op << "\n");
 
     OpBuilder builder(op);
     Location loc = op.getLoc();
-    TypedValue<tt::TensorDescType> tDesc = op.getDesc();
-    tt::TensorDescType tDescType = tDesc.getType();
-    tt::MakeTensorDescOp makeTensorDescOp = getMakeTensorDescOp(tDesc);
-
-    if (!makeTensorDescOp) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "could not find tt.make_tensor_descriptor defining: "
-                 << tDesc << "\n");
-      return failure();
-    }
-
-    LLVM_DEBUG(llvm::dbgs() << "which has tdesc: " << makeTensorDescOp << "\n");
+    tt::TensorDescType tDescType = op.getType();
 
     // Create a new block pointer if a suitable one doesn't already exist.
     SmallVector<Value> shapes, strides, offsets;
     SmallVector<int32_t> sizes;
-    for (const auto [shape, stride, offset, size] :
-         llvm::zip(makeTensorDescOp.getShape(), makeTensorDescOp.getStrides(),
-                   op.getIndices(), tDescType.getBlockType().getShape())) {
+    for (const auto [shape, stride, size] :
+         llvm::zip(op.getShape(), op.getStrides(),
+                   tDescType.getBlockType().getShape())) {
       shapes.push_back(findOrCreateCast(
           loc, shape, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
       strides.push_back(findOrCreateCast(
           loc, stride, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
-      offsets.push_back(findOrCreateCast(
-          loc, offset, builder.getIntegerType(offsetBitwidth), builder));
+      Value zero =
+          tt::intel::findOrCreateIntConstant(loc, 0, offsetBitwidth, builder);
+      offsets.push_back(zero);
path: third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
user: login: etiotto
body: The make_tensor_ptr operation being created requires offsets but the original make_tensor_descriptor operation doesn't have offsets, so this code review comment is incorrect.
created_at: 2025-05-20T17:04:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,25 +1,27 @@
 // RUN: triton-opt %s -triton-intel-tdesc-to-block-pointer  | FileCheck %s
 
-module {
path: test/Triton/Intel/TensorDescToBlockPointer/basic.mlir
user: login: whitneywhtsang
body: curious why the test case now require num-warps?
created_at: 2025-05-20T21:16:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -147,77 +113,106 @@ struct TritonIntelTensorDescToBlockPointer
       return false;
     });
 
+    auto makeTensorPtrOp = [&]() {
+      Value makeTensorPtr = builder.create<tt::MakeTensorPtrOp>(
+          loc, base, shape, strides, offsets, sizes,
+          builder.getDenseI32ArrayAttr({1, 0}));
+      return makeTensorPtr;
+    };
+
     return (it != insertPoint) ? cast<tt::MakeTensorPtrOp>(*it)
-                               : builder.createOrFold<tt::MakeTensorPtrOp>(
-                                     loc, base, shape, strides, offsets, sizes,
-                                     builder.getDenseI32ArrayAttr({1, 0}));
+                               : makeTensorPtrOp();
   }
 
-  template <typename OpTy,
-            std::enable_if_t<llvm::is_one_of<OpTy, tt::DescriptorLoadOp,
-                                             tt::DescriptorStoreOp>::value,
-                             bool> = true>
-  LogicalResult rewriteDescriptorLoadOrStoreOp(OpTy op) {
+  LogicalResult rewriteMakeTensorDescriptorOp(tt::MakeTensorDescOp op) {
     assert(op && "Expecting a valid operation");
     LLVM_DEBUG(llvm::dbgs() << "Rewriting: " << op << "\n");
 
     OpBuilder builder(op);
     Location loc = op.getLoc();
-    TypedValue<tt::TensorDescType> tDesc = op.getDesc();
-    tt::TensorDescType tDescType = tDesc.getType();
-    tt::MakeTensorDescOp makeTensorDescOp = getMakeTensorDescOp(tDesc);
-
-    if (!makeTensorDescOp) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "could not find tt.make_tensor_descriptor defining: "
-                 << tDesc << "\n");
-      return failure();
-    }
-
-    LLVM_DEBUG(llvm::dbgs() << "which has tdesc: " << makeTensorDescOp << "\n");
+    tt::TensorDescType tDescType = op.getType();
 
     // Create a new block pointer if a suitable one doesn't already exist.
     SmallVector<Value> shapes, strides, offsets;
     SmallVector<int32_t> sizes;
-    for (const auto [shape, stride, offset, size] :
-         llvm::zip(makeTensorDescOp.getShape(), makeTensorDescOp.getStrides(),
-                   op.getIndices(), tDescType.getBlockType().getShape())) {
+    for (const auto [shape, stride, size] :
+         llvm::zip(op.getShape(), op.getStrides(),
+                   tDescType.getBlockType().getShape())) {
       shapes.push_back(findOrCreateCast(
           loc, shape, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
       strides.push_back(findOrCreateCast(
           loc, stride, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
-      offsets.push_back(findOrCreateCast(
-          loc, offset, builder.getIntegerType(offsetBitwidth), builder));
+      Value zero =
+          tt::intel::findOrCreateIntConstant(loc, 0, offsetBitwidth, builder);
+      offsets.push_back(zero);
       sizes.push_back(static_cast<int32_t>(size));
     }
 
-    Value makeTensorPtrOp =
-        findOrCreateMakeTensorPtr(loc, makeTensorDescOp.getBase(), shapes,
-                                  strides, offsets, sizes, builder);
-
+    Value tensorPtr = findOrCreateMakeTensorPtr(
+        loc, op.getBase(), shapes, strides, offsets, sizes, builder);
     LLVM_DEBUG({
       llvm::dbgs() << "With:\n";
-      llvm::dbgs().indent(2) << makeTensorPtrOp << "\n";
+      llvm::dbgs().indent(2) << tensorPtr << "\n";
     });
 
+    op.replaceAllUsesWith(tensorPtr);
+    cleanUp.insert(op);
+
+    for (Operation *user : tensorPtr.getUsers()) {
+      if (auto forOp = dyn_cast<scf::ForOp>(user)) {
+        for (auto it :
+             llvm::zip(forOp.getInitArgs(), forOp.getRegionIterArgs(),
+                       forOp.getResults(), forOp.getYieldedValues())) {
+          Value initArg = std::get<0>(it), rgnInitArg = std::get<1>(it),
+                loopRes = std::get<2>(it), yieldVal = std::get<3>(it);
path: third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
user: login: whitneywhtsang
body: ```suggestion
        for (auto [initArg, rgnInitArg, loopRes, yieldVal] :
             llvm::zip(forOp.getInitArgs(), forOp.getRegionIterArgs(),
                       forOp.getResults(), forOp.getYieldedValues())) {
```
created_at: 2025-05-20T21:24:56Z
author_association: CONTRIBUTOR
path: third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
user: login: whitneywhtsang
body: should we add a lit test in `test/TritonIntelGPU/combine.mlir` for this change?
created_at: 2025-05-20T21:26:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -692,6 +692,26 @@ void LayoutPropagation::rewriteAssertOp(AssertOp assertOp) {
   assertOp->setOperand(0, newOperand);
 }
 
+// Recursively update the operands in a chain of AdvanceOps, after setting the
+// pointer operand of the first one.
+static void updateAdvanceOpChain(AdvanceOp advanceOp, Value makeTensorPtrOp,
+                                 Value dataToStore) {
+  OpBuilder rewriter(advanceOp);
+  auto newAdvanceOp =
+      rewriter.create<AdvanceOp>(advanceOp.getLoc(), makeTensorPtrOp.getType(),
+                                 makeTensorPtrOp, advanceOp.getOffsets());
+
+  SmallVector<Operation *> advanceOpUsers(advanceOp->getUsers());
+  for (Operation *user : advanceOpUsers) {
+    if (auto storeOp = dyn_cast<StoreOp>(user)) {
+      storeOp.setOperand(0, newAdvanceOp);
+      storeOp.setOperand(1, dataToStore);
+    } else if (auto advanceOp = dyn_cast<AdvanceOp>(user)) {
+      updateAdvanceOpChain(advanceOp, makeTensorPtrOp, dataToStore);
+    }
path: third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
user: login: whitneywhtsang
body: should we assert that there is no other kind of users?
created_at: 2025-05-20T21:30:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -758,21 +781,29 @@ bool LayoutPropagation::rewriteStoreOp(StoreOp storeOp) {
     encoding = convertOpSrcType.getEncoding();
   }
 
-  // We create a new MakeTensorPtrOp with the new data type.
+  // Create a new MakeTensorPtrOp with the new layout.
   OpBuilder rewriter(makeTensorPtrOp);
-  Value newStorePtr = rewriter.create<MakeTensorPtrOp>(
+  Value newMakeTensorPtrOp = rewriter.create<MakeTensorPtrOp>(
       makeTensorPtrOp.getLoc(), newPtrType, makeTensorPtrOp.getBase(),
       makeTensorPtrOp.getShape(), makeTensorPtrOp.getStrides(),
-      makeTensorPtrOp.getOffsets(), rewriter.getDenseI32ArrayAttr({1, 0}));
+      makeTensorPtrOp.getOffsets(), makeTensorPtrOp.getOrderAttr());
+
+  // Update the store operation with the new layout.
+  SmallVector<Operation *> makeTensorPtrOpUsers(makeTensorPtrOp->getUsers());
+  Value dataToStore = getValueAs(value, encoding);
+  Block *storeBB = storeOp->getBlock();
+  for (Operation *user : makeTensorPtrOpUsers) {
+    Block *userBB = user->getBlock();
+    if (storeBB != userBB)
path: third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
user: login: whitneywhtsang
body: why only need to update when the two blocks are the same?
created_at: 2025-05-20T21:35:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,25 +1,27 @@
 // RUN: triton-opt %s -triton-intel-tdesc-to-block-pointer  | FileCheck %s
 
-module {
path: test/Triton/Intel/TensorDescToBlockPointer/basic.mlir
user: login: etiotto
body: It is superfluous, removing it.
created_at: 2025-05-22T15:32:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -147,77 +113,106 @@ struct TritonIntelTensorDescToBlockPointer
       return false;
     });
 
+    auto makeTensorPtrOp = [&]() {
+      Value makeTensorPtr = builder.create<tt::MakeTensorPtrOp>(
+          loc, base, shape, strides, offsets, sizes,
+          builder.getDenseI32ArrayAttr({1, 0}));
+      return makeTensorPtr;
+    };
+
     return (it != insertPoint) ? cast<tt::MakeTensorPtrOp>(*it)
-                               : builder.createOrFold<tt::MakeTensorPtrOp>(
-                                     loc, base, shape, strides, offsets, sizes,
-                                     builder.getDenseI32ArrayAttr({1, 0}));
+                               : makeTensorPtrOp();
   }
 
-  template <typename OpTy,
-            std::enable_if_t<llvm::is_one_of<OpTy, tt::DescriptorLoadOp,
-                                             tt::DescriptorStoreOp>::value,
-                             bool> = true>
-  LogicalResult rewriteDescriptorLoadOrStoreOp(OpTy op) {
+  LogicalResult rewriteMakeTensorDescriptorOp(tt::MakeTensorDescOp op) {
     assert(op && "Expecting a valid operation");
     LLVM_DEBUG(llvm::dbgs() << "Rewriting: " << op << "\n");
 
     OpBuilder builder(op);
     Location loc = op.getLoc();
-    TypedValue<tt::TensorDescType> tDesc = op.getDesc();
-    tt::TensorDescType tDescType = tDesc.getType();
-    tt::MakeTensorDescOp makeTensorDescOp = getMakeTensorDescOp(tDesc);
-
-    if (!makeTensorDescOp) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "could not find tt.make_tensor_descriptor defining: "
-                 << tDesc << "\n");
-      return failure();
-    }
-
-    LLVM_DEBUG(llvm::dbgs() << "which has tdesc: " << makeTensorDescOp << "\n");
+    tt::TensorDescType tDescType = op.getType();
 
     // Create a new block pointer if a suitable one doesn't already exist.
     SmallVector<Value> shapes, strides, offsets;
     SmallVector<int32_t> sizes;
-    for (const auto [shape, stride, offset, size] :
-         llvm::zip(makeTensorDescOp.getShape(), makeTensorDescOp.getStrides(),
-                   op.getIndices(), tDescType.getBlockType().getShape())) {
+    for (const auto [shape, stride, size] :
+         llvm::zip(op.getShape(), op.getStrides(),
+                   tDescType.getBlockType().getShape())) {
       shapes.push_back(findOrCreateCast(
           loc, shape, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
       strides.push_back(findOrCreateCast(
           loc, stride, builder.getIntegerType(shapeAndStridesBitwidth),
           builder));
-      offsets.push_back(findOrCreateCast(
-          loc, offset, builder.getIntegerType(offsetBitwidth), builder));
+      Value zero =
+          tt::intel::findOrCreateIntConstant(loc, 0, offsetBitwidth, builder);
+      offsets.push_back(zero);
       sizes.push_back(static_cast<int32_t>(size));
     }
 
-    Value makeTensorPtrOp =
-        findOrCreateMakeTensorPtr(loc, makeTensorDescOp.getBase(), shapes,
-                                  strides, offsets, sizes, builder);
-
+    Value tensorPtr = findOrCreateMakeTensorPtr(
+        loc, op.getBase(), shapes, strides, offsets, sizes, builder);
     LLVM_DEBUG({
       llvm::dbgs() << "With:\n";
-      llvm::dbgs().indent(2) << makeTensorPtrOp << "\n";
+      llvm::dbgs().indent(2) << tensorPtr << "\n";
     });
 
+    op.replaceAllUsesWith(tensorPtr);
+    cleanUp.insert(op);
+
+    for (Operation *user : tensorPtr.getUsers()) {
+      if (auto forOp = dyn_cast<scf::ForOp>(user)) {
+        for (auto it :
+             llvm::zip(forOp.getInitArgs(), forOp.getRegionIterArgs(),
+                       forOp.getResults(), forOp.getYieldedValues())) {
+          Value initArg = std::get<0>(it), rgnInitArg = std::get<1>(it),
+                loopRes = std::get<2>(it), yieldVal = std::get<3>(it);
path: third_party/intel/lib/Dialect/Triton/Transforms/TensorDescToBlockPointer.cpp
user: login: etiotto
body: More concise, changing it. Thanks!
created_at: 2025-05-22T15:34:48Z
author_association: CONTRIBUTOR
path: third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
user: login: etiotto
body: adding one
created_at: 2025-05-22T16:46:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -758,21 +781,29 @@ bool LayoutPropagation::rewriteStoreOp(StoreOp storeOp) {
     encoding = convertOpSrcType.getEncoding();
   }
 
-  // We create a new MakeTensorPtrOp with the new data type.
+  // Create a new MakeTensorPtrOp with the new layout.
   OpBuilder rewriter(makeTensorPtrOp);
-  Value newStorePtr = rewriter.create<MakeTensorPtrOp>(
+  Value newMakeTensorPtrOp = rewriter.create<MakeTensorPtrOp>(
       makeTensorPtrOp.getLoc(), newPtrType, makeTensorPtrOp.getBase(),
       makeTensorPtrOp.getShape(), makeTensorPtrOp.getStrides(),
-      makeTensorPtrOp.getOffsets(), rewriter.getDenseI32ArrayAttr({1, 0}));
+      makeTensorPtrOp.getOffsets(), makeTensorPtrOp.getOrderAttr());
+
+  // Update the store operation with the new layout.
+  SmallVector<Operation *> makeTensorPtrOpUsers(makeTensorPtrOp->getUsers());
+  Value dataToStore = getValueAs(value, encoding);
+  Block *storeBB = storeOp->getBlock();
+  for (Operation *user : makeTensorPtrOpUsers) {
+    Block *userBB = user->getBlock();
+    if (storeBB != userBB)
path: third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
user: login: etiotto
body: Spit PR, please add comment in #4277 
created_at: 2025-05-22T18:28:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -692,6 +692,26 @@ void LayoutPropagation::rewriteAssertOp(AssertOp assertOp) {
   assertOp->setOperand(0, newOperand);
 }
 
+// Recursively update the operands in a chain of AdvanceOps, after setting the
+// pointer operand of the first one.
+static void updateAdvanceOpChain(AdvanceOp advanceOp, Value makeTensorPtrOp,
+                                 Value dataToStore) {
+  OpBuilder rewriter(advanceOp);
+  auto newAdvanceOp =
+      rewriter.create<AdvanceOp>(advanceOp.getLoc(), makeTensorPtrOp.getType(),
+                                 makeTensorPtrOp, advanceOp.getOffsets());
+
+  SmallVector<Operation *> advanceOpUsers(advanceOp->getUsers());
+  for (Operation *user : advanceOpUsers) {
+    if (auto storeOp = dyn_cast<StoreOp>(user)) {
+      storeOp.setOperand(0, newAdvanceOp);
+      storeOp.setOperand(1, dataToStore);
+    } else if (auto advanceOp = dyn_cast<AdvanceOp>(user)) {
+      updateAdvanceOpChain(advanceOp, makeTensorPtrOp, dataToStore);
+    }
path: third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
user: login: etiotto
body: Spit PR, please add comment in #4277 
created_at: 2025-05-22T18:28:51Z
author_association: CONTRIBUTOR

PR comments:
user: login: kurapov-peter
created_at: 2025-05-14T11:06:40Z
body: BTW FYI, the existing TMA lowering expects the tensor descriptor to always have a layout.
author_association: CONTRIBUTOR
user: login: etiotto
created_at: 2025-05-15T21:09:40Z
body: > The code LGTM, but I'm not sure to understand why we need this PR if the conversion from `tensor_descriptor` to ` block_pointer` is one of the first passes to run at the `ttir` level, how can we have `tensor_descriptor` ops with different layouts given the encodings are assigned to ops when lowering from the `ttir` level to the `ttgir` level. As I understand it, we are not supposed to have `tensor_descriptor` ops after translating to `ttgir` dialect, no?

I have updated the code quite a bit to change the loginc used to do the translation. The pass is now simpler and able to handle tesnor descriptors created in control flow (i.e. a branch) inside a loop.
author_association: CONTRIBUTOR
user: login: etiotto
created_at: 2025-05-20T16:24:18Z
body: No performance degradation in Triton benchmarks: https://github.com/intel-sandbox/applications.python.intel-xpu-backend-for-triton.infrastructure/actions/runs/15142551794/job/42570177577
author_association: CONTRIBUTOR
user: login: etiotto
created_at: 2025-05-22T18:19:51Z
body: @whitneywhtsang I have split this pull request and move the part that deals with removing layout conversion for store operation that use a block ptr updated by a tt.advance operation in PR #4277 
author_association: CONTRIBUTOR
user: login: whitneywhtsang
created_at: 2025-05-22T18:31:12Z
body: > @whitneywhtsang I have split this pull request and move the part that deals with removing layout conversion for store operation that use a block ptr updated by a tt.advance operation in PR #4277

Thanks, this part of the code LGTM.
author_association: CONTRIBUTOR
user: login: etiotto
created_at: 2025-05-28T12:54:35Z
body: Benchmark results: https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/15285791268
No degradations detected.
author_association: CONTRIBUTOR