Update per comment.

Author: Xiang Li

lib/AnalysisStructured/PtrAnalysis.cpp

@@ -53,22 +53,23 @@ static Value applyUnstructuredMask(Operation *op, Value ptr,
   }
 
   auto [dim, unstructuredMask] = masks[0];
-  if (auto scatterPtr =
+  if (auto gatherScatterPtr =
           ptr.getDefiningOp<tts::MakeGatherScatterTensorPtrOp>()) {
-    if (dim != scatterPtr.getGatherScatterDim()) {
+    if (dim != gatherScatterPtr.getGatherScatterDim()) {
       op->emitRemark("MaskAnalysis failed for unstructured mask dim not equal "
                      "gather scatter dim");
       return nullptr;
     }
 
-    ptr = builder
-              .create<tts::MakeGatherScatterTensorPtrOp>(
-                  loc, scatterPtr.getBase(),
-                  scatterPtr.getGatherScatterOffset(), unstructuredMask,
-                  scatterPtr.getGatherScatterDim(), scatterPtr.getSizes(),
-                  scatterPtr.getMixedStrides(), scatterPtr.getMixedOffsets())
-              .getResult();
-
+    ptr =
+        builder
+            .create<tts::MakeGatherScatterTensorPtrOp>(
+                loc, gatherScatterPtr.getBase(),
+                gatherScatterPtr.getGatherScatterOffset(), unstructuredMask,
+                gatherScatterPtr.getGatherScatterDim(),
+                gatherScatterPtr.getSizes(), gatherScatterPtr.getMixedStrides(),
+                gatherScatterPtr.getMixedOffsets())
+            .getResult();
   } else if (auto structuredPtr = ptr.getDefiningOp<tts::MakeTensorPtrOp>()) {
     auto ofrToI32Value = [&](OpFoldResult ofr) {
       Value v = dyn_cast<Value>(ofr);

test/Conversion/TritonToStructured/unstructured_mask_2d_kernel.mlir

@@ -1,6 +1,8 @@
 // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s
 
-// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from structured ptr with unstructured mask.
+// Make sure make_gather_scatter_tptr with unstructuredmask generate correctly from structured ptr with unstructured mask.
+// The load is structured ptr, with unstructured mask on dim 0.
+// The store is structured ptr, with unstructured mask on dim 1.
 
 // CHECK-LABEL:   tt.func public @generic_mask_2d_kernel(
 // CHECK-SAME:                                           %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32},

test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_load_kernel.mlir

@@ -1,6 +1,7 @@
 // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s
 
-// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from row-structured ptr with unstructured mask.
+// Make sure make_gather_scatter_tptr with unstructuredmask generate correctly from row-structured ptr with unstructured mask.
+// The load is unstructured ptr on dim 0 and unstructured mask on dim 0.
 
 // CHECK-LABEL:   tt.func public @generic_mask_2d_non_continuous_load_kernel(
 // CHECK-SAME:                                                               %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32},

test/Conversion/TritonToStructured/unstructured_mask_2d_non_continuous_store_kernel.mlir

@@ -1,6 +1,7 @@
 // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s
 
-// Make sure make_gather_scatter_tptr with unsturctured mask generate correctly from column-structured ptr with unstructured mask.
+// Make sure make_gather_scatter_tptr with unstructuredmask generate correctly from column-structured ptr with unstructured mask.
+// The store is unstructured ptr on dim 1 and unstructured mask on dim 1.
 
 // CHECK-LABEL:   tt.func public @generic_mask_2d_non_continuous_store_kernel(
 // CHECK-SAME:                                                                %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32},

test/Conversion/TritonToStructured/unstructured_mask_3d_kernel.mlir

@@ -1,6 +1,8 @@
 // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s
 
 // Make sure make_gather_scatter_tptr with generic mask generate correctly from structured ptr with unstructured mask.
+// The load is structured ptr, with unstructured mask on dim 1.
+// The store is structured ptr, with unstructured mask on dim 2.
 
 // CHECK-LABEL:   tt.func public @generic_mask_3d_kernel(
 // CHECK-SAME:                                           %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32},

test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_load_kernel.mlir

@@ -1,6 +1,7 @@
 // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s
 
 // Make sure make_gather_scatter_tptr with generic mask generate correctly from row-structured ptr with unstructured mask.
+// The load is unstructured ptr on dim 1 and unstructured mask on dim 1.
 
 // CHECK-LABEL:   tt.func public @generic_mask_3d_non_continuous_load_kernel(
 // CHECK-SAME:                                                                %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !tt.ptr<f32> {tt.divisibility = 16 : i32},

test/Conversion/TritonToStructured/unstructured_mask_3d_non_continuous_store_kernel.mlir

@@ -1,6 +1,7 @@
 // RUN: triton-shared-opt --triton-to-structured --remove-dead-values --cse --canonicalize %s | FileCheck %s
 
 // Make sure make_gather_scatter_tptr with generic mask generate correctly from column-structured ptr with unstructured mask.
+// The store is unstructured ptr on dim 2 and unstructured mask on dim 2.
 
 
 // CHECK-LABEL:   tt.func public @generic_mask_3d_non_continuous_store_kernel(