[Mlir-commits] [mlir] f18e14d - [mlir][Linalg] NFC - Add transform pack builder

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2023-01-24T02:26:15-08:00
New Revision: f18e14d863363d93a587bf6f25096d29d08f25b3

URL: https://github.com/llvm/llvm-project/commit/f18e14d863363d93a587bf6f25096d29d08f25b3
DIFF: https://github.com/llvm/llvm-project/commit/f18e14d863363d93a587bf6f25096d29d08f25b3.diff

LOG: [mlir][Linalg] NFC - Add transform pack builder

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
    mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
    mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 991aa041b9531..804ce5b370eb1 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -366,6 +366,7 @@ def MultiTileSizesOp : Op<Transform_Dialect, "structured.multitile_sizes",
 //===----------------------------------------------------------------------===//
 // PackOp
 //===----------------------------------------------------------------------===//
+
 def PackOp : Op<Transform_Dialect, "structured.pack", [
                 DeclareOpInterfaceMethods<TransformOpInterface>,
                 DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,]> {
@@ -441,6 +442,11 @@ def PackOp : Op<Transform_Dialect, "structured.pack", [
     `:` functional-type($target, results)
   }];
 
+  let builders = [
+    OpBuilder<(ins "Value":$target,
+                   "ArrayRef<OpFoldResult>":$mixedPackedSizes)>
+  ];
+
   let extraClassDeclaration = [{
     ::llvm::SmallVector<::mlir::OpFoldResult> getMixedPackedSizes();
   }];

diff  --git a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
index 2dc93c1fb7ef3..4bb6700dbca29 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
@@ -225,7 +225,7 @@ def GetProducerOfOperand : TransformDialectOp<"get_producer_of_operand",
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        I64Attr:$operand_number);
-  let results = (outs TransformHandleTypeInterface:$parent);
+  let results = (outs TransformHandleTypeInterface:$producer);
   let assemblyFormat = "$target `[` $operand_number `]` attr-dict `:` "
                        "functional-type(operands, results)";
 }

diff  --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 7f3ac242d7d9f..3df23733452b2 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -874,6 +874,25 @@ LogicalResult transform::MultiTileSizesOp::verify() {
 // PackOp
 //===---------------------------------------------------------------------===//
 
+void transform::PackOp::build(OpBuilder &builder, OperationState &result,
+                              Value target,
+                              ArrayRef<OpFoldResult> mixedPackedSizes) {
+  SmallVector<int64_t> staticPackedSizes;
+  SmallVector<Value> dynamicPackedSizes;
+  dispatchIndexOpFoldResults(mixedPackedSizes, dynamicPackedSizes,
+                             staticPackedSizes);
+  // Call the default builder which sets up the proper operands segment sizes
+  // attributes for multiple variadic operands. In the absence of this, horrible
+  // bugs ensue.
+  Type linalgOpHType = transform::OperationType::get(
+      builder.getContext(), linalg::GenericOp::getOperationName());
+  build(builder, result,
+        /*resultType=*/linalgOpHType,
+        /*target=*/target,
+        /*dynamic_sizes=*/dynamicPackedSizes,
+        /*static_sizes=*/builder.getDenseI64ArrayAttr(staticPackedSizes));
+}
+
 SmallVector<OpFoldResult> transform::PackOp::getMixedPackedSizes() {
   Builder b(getContext());
   return getMixedValues(getStaticPackedSizes(), getPackedSizes(), b);