[IR] Update names of decomposed tensor descriptor args (triton-lang#9587)

Tensor descriptor function arguments are decomposed into
    tt.ptr, shape, stride, and some flags. This pr updates each
    new parameter with a meaningful name, aligned with codegen'd
    tensordesc naming.

Author: sjw36

include/triton/Dialect/Triton/Transforms/FunctionTypeConversion.h

@@ -1,9 +1,65 @@
 #ifndef TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_
 #define TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir::triton {
 
+/**
+ * @brief Rename the aggregated function arguments that were generated by the
+ * type converter.
+ *
+ * @note The callback should return a std::optional<LogicalResult>. If the
+ * callback returns a std::nullopt, the argument is not renamed and the
+ * conversion continues. If the callback returns a LogicalResult, the out_suffix
+ * list is populated with the suffixes for each aggregated argument.
+ *
+ * @param delimiter The delimiter to use between the base name and the suffixes.
+ */
+class FuncArgRenamer {
+public:
+  FuncArgRenamer(const char *delimiter = ".") : delimiter(delimiter) {}
+  template <typename FnT, typename T = typename llvm::function_traits<
+                              std::decay_t<FnT>>::template arg_t<0>>
+  void addRenamer(FnT &&callback) {
+    renamers.emplace_back(wrapCallback<T>(std::forward<FnT>(callback)));
+  }
+  LogicalResult apply(Type type, FunctionOpInterface funcOp, int index,
+                      TypeConverter::SignatureConversion &conversion) const;
+
+private:
+  using RenamerCallbackFn = std::function<std::optional<LogicalResult>(
+      Type, llvm::SmallVectorImpl<std::string> &)>;
+
+  /**
+   * @brief Wraps a callback of form `std::optional<LogicalResult>(T,
+   * llvm::SmallVectorImpl<std::string> &)` into a RenamerCallbackFn.
+   *
+   * @tparam T The type of the argument.
+   * @tparam FnT The type of the callback.
+   * @param callback The callback to wrap.
+   * @return A RenamerCallbackFn.
+   */
+  template <typename T, typename FnT>
+  std::enable_if_t<
+      std::is_invocable_v<FnT, T, llvm::SmallVectorImpl<std::string> &> &&
+          std::is_base_of_v<Type, T>,
+      RenamerCallbackFn>
+  wrapCallback(FnT &&callback) const {
+    return [callback = std::forward<FnT>(callback)](
+               Type type, llvm::SmallVectorImpl<std::string> &out_suffix)
+               -> std::optional<LogicalResult> {
+      if (auto t = dyn_cast<T>(type)) {
+        return callback(t, out_suffix);
+      }
+      return std::nullopt;
+    };
+  }
+
+  llvm::SmallVector<RenamerCallbackFn> renamers;
+  const char *delimiter;
+};
+
 /**
  * @brief Provides helper patterns for converting triton function operations
  * using a type converter.
@@ -12,6 +68,7 @@ namespace mlir::triton {
  * tt.call and tt.return.
  */
 void populateFunctionTypeConversions(const TypeConverter &converter,
+                                     const FuncArgRenamer &renamer,
                                      RewritePatternSet &patterns);
 
 } // namespace mlir::triton

lib/Dialect/Triton/Transforms/FunctionTypeConversion.cpp

@@ -11,6 +11,42 @@
 
 namespace mlir::triton {
 
+LogicalResult
+FuncArgRenamer::apply(Type type, FunctionOpInterface funcOp, int index,
+                      TypeConverter::SignatureConversion &conversion) const {
+  auto mapping = conversion.getInputMapping(index);
+  if (!mapping)
+    return success();
+
+  for (auto &renamer : llvm::reverse(renamers)) {
+    llvm::SmallVector<std::string, 8> out_suffix;
+    if (std::optional<LogicalResult> result = renamer(type, out_suffix)) {
+      if (failed(*result)) {
+        return failure();
+      }
+      int newIndex = mapping->inputNo;
+      auto loc = funcOp.getArgument(newIndex).getLoc();
+      std::string baseName;
+      if (isa<NameLoc>(loc)) {
+        baseName = cast<NameLoc>(loc).getName().getValue();
+      } else {
+        baseName = "arg_" + std::to_string(index);
+      }
+      assert(out_suffix.size() == mapping->size);
+      for (auto [i, suffix] : llvm::enumerate(out_suffix)) {
+        if (suffix.empty())
+          continue;
+        auto newLoc = NameLoc::get(
+            StringAttr::get(funcOp.getContext(), baseName + delimiter + suffix),
+            loc);
+        funcOp.getArgument(newIndex + i).setLoc(newLoc);
+      }
+      return success(); // early return
+    }
+  }
+  return success();
+}
+
 namespace {
 
 SmallVector<Value> flattenValues(ArrayRef<ValueRange> values) {
@@ -102,6 +138,7 @@ convertFuncOpAttrs(FunctionOpInterface funcOp,
 
 LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
                                  const TypeConverter &typeConverter,
+                                 const FuncArgRenamer &renamer,
                                  ConversionPatternRewriter &rewriter) {
   FunctionType type = dyn_cast<FunctionType>(funcOp.getFunctionType());
   if (!type)
@@ -129,6 +166,13 @@ LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
     }
   });
 
+  // Apply the renamer to the function signature.
+  for (auto [i, input] : llvm::enumerate(type.getInputs())) {
+    if (failed(renamer.apply(input, funcOp, i, result))) {
+      return failure();
+    }
+  }
+
   return success();
 }
 
@@ -139,24 +183,30 @@ struct FunctionOpInterfaceSignatureConversion : public ConversionPattern {
   FunctionOpInterfaceSignatureConversion(StringRef functionLikeOpName,
                                          MLIRContext *ctx,
                                          const TypeConverter &converter,
+                                         const FuncArgRenamer &renamer,
                                          PatternBenefit benefit = 1)
-      : ConversionPattern(converter, functionLikeOpName, benefit, ctx) {}
+      : ConversionPattern(converter, functionLikeOpName, benefit, ctx),
+        renamer(renamer) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> /*operands*/,
                   ConversionPatternRewriter &rewriter) const override {
     FunctionOpInterface funcOp = cast<FunctionOpInterface>(op);
-    return convertFuncOpTypes(funcOp, *typeConverter, rewriter);
+    return convertFuncOpTypes(funcOp, *typeConverter, renamer, rewriter);
   }
+
+private:
+  const FuncArgRenamer &renamer;
 };
 
 } // namespace
 
 void populateFunctionTypeConversions(const TypeConverter &converter,
+                                     const FuncArgRenamer &renamer,
                                      RewritePatternSet &patterns) {
   auto context = patterns.getContext();
   patterns.add<FunctionOpInterfaceSignatureConversion>(
-      triton::FuncOp::getOperationName(), context, converter);
+      triton::FuncOp::getOperationName(), context, converter, renamer);
   patterns.add<CallOpConversion, ReturnOpConversion>(converter, context);
 }
 

lib/Dialect/Triton/Transforms/RewriteTensorDescriptorToPointer.cpp

@@ -580,11 +580,28 @@ class TritonRewriteTensorDescriptorToPointerPass
       return mlir::success();
     });
 
+    FuncArgRenamer renamer(".");
+    renamer.addRenamer([](mlir::triton::TensorDescType type,
+                          llvm::SmallVectorImpl<std::string> &out_suffix) {
+      auto tensorType = type.getSignlessBlockType();
+      int dims = tensorType.getRank();
+      out_suffix.push_back("");
+      for (int i = 0; i < dims; i++) {
+        out_suffix.push_back("shape." + std::to_string(i));
+      }
+      for (int i = 0; i < dims; i++) {
+        out_suffix.push_back("stride." + std::to_string(i));
+      }
+      out_suffix.push_back("padding");
+      out_suffix.push_back("roundF32ToTF32");
+      return success();
+    });
+
     mlir::RewritePatternSet patterns(op->getContext());
 
     // Populate conversion patterns to handle loops, function calls, and arith
     // ops.
-    triton::populateFunctionTypeConversions(converter, patterns);
+    triton::populateFunctionTypeConversions(converter, renamer, patterns);
     mlir::scf::populateSCFStructuralTypeConversions(converter, patterns);
     triton::populateArithTypeConversions(converter, patterns);
 

test/Triton/rewrite-tensor-descriptor-to-pointer.mlir

@@ -1,4 +1,4 @@
-// RUN: triton-opt %s --triton-rewrite-tensor-descriptor-to-pointer --canonicalize --cse --split-input-file | FileCheck %s --implicit-check-not \!tt.tensordesc
+// RUN: triton-opt %s --triton-rewrite-tensor-descriptor-to-pointer --canonicalize --cse --mlir-print-debuginfo --split-input-file | FileCheck %s --implicit-check-not \!tt.tensordesc
 
 module {
   tt.func public @load(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: i32, %arg2: i32) -> (tensor<128x128xf32>) {
@@ -109,8 +109,9 @@ module {
 
 // -----
 
+#loc2 = loc("rewrite-tensor-descriptor-to-pointer.mlir":147:28)
 module {
-  tt.func public @callee(%tensordesc: !tt.tensordesc<tensor<128x128xf32>>) -> !tt.tensordesc<tensor<128x128xf32>> {
+  tt.func public @callee(%tensordesc: !tt.tensordesc<tensor<128x128xf32>> loc("tensordesc"(#loc2))) -> !tt.tensordesc<tensor<128x128xf32>> {
     tt.return %tensordesc : !tt.tensordesc<tensor<128x128xf32>>
   }
 
@@ -126,12 +127,19 @@ module {
 
 // CHECK-LABEL: @callee
 // CHECK-SAME: %[[PTR:[^:]*]]
+// CHECK-SAME: loc("tensordesc"(#loc{{[^,]*}}))
 // CHECK-SAME: %[[SHAPE0:[^:]*]]
+// CHECK-SAME: loc("tensordesc.shape.0"(#loc{{[^,]*}}))
 // CHECK-SAME: %[[SHAPE1:[^:]*]]
+// CHECK-SAME: loc("tensordesc.shape.1"(#loc{{[^,]*}}))
 // CHECK-SAME: %[[STRIDE0:[^:]*]]
+// CHECK-SAME: loc("tensordesc.stride.0"(#loc{{[^,]*}}))
 // CHECK-SAME: %[[STRIDE1:[^:]*]]
+// CHECK-SAME: loc("tensordesc.stride.1"(#loc{{[^,]*}}))
 // CHECK-SAME: %[[PAD:[^:]*]]
+// CHECK-SAME: loc("tensordesc.padding"(#loc{{[^,]*}}))
 // CHECK-SAME: %[[ROUND:[^:]*]]
+// CHECK-SAME: loc("tensordesc.roundF32ToTF32"(#loc{{[^,]*}}))
 // CHECK-NEXT: tt.return %[[PTR]], %[[SHAPE0]], %[[SHAPE1]], %[[STRIDE0]], %[[STRIDE1]], %[[PAD]], %[[ROUND]]
 
 // CHECK-LABEL: @caller
@@ -150,4 +158,4 @@ module {
 }
 
 // CHECK-LABEL: @arg_attr
-// CHECK-SAME: %arg7: i32 {tt.divisibility = 16 : i32}) {
+// CHECK-SAME: %arg7: i32 {tt.divisibility = 16 : i32} loc({{.*}})) {