[mlir][Transform] Add support for dynamically unpacking tile_sizes / num_threads in tile_to_foreach_thread

[nicolasvasilache]

Summary

This commit adds automatic unpacking of Value's of type pdl::OperationType to the underlying single-result OpResult.
This allows mixing single-value, attribute and multi-value pdl::Operation tile sizes and num threads to TileToForeachThreadOp.

[nicolasvasilache]

Update.

[nicolasvasilache]

Update.

[nicolasvasilache]

Update commit message.

[nicolasvasilache]

Rebase.

[springerm]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

799

to

799–800

I would rephrase this:

`tile_sizes` and `num_threads` are variadic. Each tile size/number of threads can be an index attribute or a transform handle that is mapped to exactly one payload op with exactly one index result.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

1330–1333

Why do you change this? The name of the attr in the .td file is mapping.
If you keep this change, also update the builder decls in the .td file.

1360

May need updating

1418

May need updating

1436

rename variable

1438

No loop needed, because the must be exactly one payload op. Raise an error if a different number of payload ops are mapped.

1483–1484

Nit: If you add this, you skip the error checking inside of unpackPDLOperations which could be useful for debugging faulty IR.

1546–1551

Why is this needed? Is this something that we should do with every transform op?

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

799–800

thanks!

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

1360

I switched back to mapping everywhere, is that the updating you meant or something else ?

1438

No, that is the whole point of the PR, the handle may point to multiple payload ops that we need to unpack.

[0, %three_tile_sizes, %c42] will unpack to 5 tiles sizes total.

1483–1484

Propagating the empty state properly through the whole transform IR is much more important as it can lead to hard to track issues when used at scale.

If this could be systematized, this would be great but also not on my short-term priority list.

1546–1551

Yes, we should always set the expected handles in the success and silenceable failure cases otherwise hard to track bugs ensue (at the very least in the "Suppress" mode).

A systematic mechanism would be very useful but not on my short term list

[nicolasvasilache]

Address comments.

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
 
     It is the user's responsibility to ensure that `num_threads/tile_sizes` is
     a valid tiling specification (i.e. that only tiles parallel dimensions, 
     e.g. in the Linalg case).
 
     If non-empty, the `mapping` is added as an attribute to the
     resulting `scf.foreach_thread`.
 
+    A `packed_sizes` unit attribute may be specified that controls whether the
+    `num_threads` or `tile_sizes` specification refers to a handle that contains
+    multiple PDL_Operation values that are unpacked dynamically.
+
     #### Return modes
 
     This operation ignores ops that do not implement the TilingInterface and
     drops them in the return.
 
     If all the operations referred to by the `target` PDLOperation tile
     successfully, the transform succeeds.
     Otherwise the transform silently fails.
let description = [{
     ```
 
     #### Example using `tile_sizes`
 
     ```
     %0 = pdl_match @match_matmul in %arg1
     %sz = pdl_match @match_size_op in %arg1
     %3:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [0, %sz, 20]
+
+    #### Example using `tile_sizes` and `packed_sizes`
+
+    ```
+    %0 = pdl_match @match_matmul in %arg1
+    %sz = pdl_match @match_multiple_size_ops in %arg1
+    %3:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [%sz] {packed_sizes}
     ```
   }];
 
   let arguments = (ins PDL_Operation:$target,
                    Variadic<PDL_Operation>:$num_threads,
                    Variadic<PDL_Operation>:$tile_sizes,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$static_num_threads,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$static_tile_sizes,
                    OptionalAttr<DeviceMappingArrayAttr>:$mapping);
   let results = (outs PDL_Operation:$foreach_thread_op,
                       PDL_Operation:$tiled_op);
 
   let builders = [
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<int64_t>":$staticTileSizes,
                    CArg<"::mlir::transform::TileSizesSpec", 
                         "::mlir::transform::TileSizesSpec()">,
                    CArg<"ArrayAttr", "{}">:$mapping)>,
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedTileSizes,
                    CArg<"::mlir::transform::TileSizesSpec", 
                         "::mlir::transform::TileSizesSpec()">,
                    CArg<"ArrayAttr", "{}">:$mapping)>,
     OpBuilder<(ins "Value":$target,
+                   "Value":$packedTileSizesHandle,
+                   CArg<"::mlir::transform::TileSizesSpec",
+                        "::mlir::transform::TileSizesSpec()">,
+                   CArg<"ArrayAttr", "{}">:$mapping)>,
+    OpBuilder<(ins "Value":$target,
                    "ArrayRef<int64_t>":$staticNumThreads,
                    CArg<"::mlir::transform::NumThreadsSpec", 
                         "::mlir::transform::NumThreadsSpec()">,
                    CArg<"ArrayAttr", "{}">:$mapping)>,
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedNumThreads,
                    CArg<"::mlir::transform::NumThreadsSpec", 
                         "::mlir::transform::NumThreadsSpec()">,
                    CArg<"ArrayAttr", "{}">:$mapping)>,
+    OpBuilder<(ins "Value":$target,
+                   "Value":$packedNumThreadsHandle,
+                   CArg<"::mlir::transform::NumThreadsSpec", 
+                        "::mlir::transform::NumThreadsSpec()">,
+                   CArg<"ArrayAttr", "{}">:$mapping)>
   ];
 
   let assemblyFormat = [{
     $target oilist(
         `num_threads` custom<DynamicIndexList>($num_threads,
                                                $static_num_threads,
                                                "ShapedType::kDynamicSize") |
          `tile_sizes` custom<DynamicIndexList>($tile_sizes,

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::transform;

 // TileToForeachThreadOp
 //===----------------------------------------------------------------------===//
 
 void transform::TileToForeachThreadOp::build(OpBuilder &builder,
                                              OperationState &result,
                                              Value target,
                                              ArrayRef<int64_t> staticTileSizes,
                                              transform::TileSizesSpec,
-                                             ArrayAttr mapping) {
-  return build(builder, result, target,
+                                             ArrayAttr mappingAttr) {
+  return build(builder, result,
+               /*target=*/target,
+               /*mixedTileSizes=*/
                getAsOpFoldResult(builder.getI64ArrayAttr(staticTileSizes)),
-               TileSizesSpec(), mapping);
+               /*_=*/TileSizesSpec(),
+               /*mapping=*/mappingAttr);
 }
 
 void transform::TileToForeachThreadOp::build(
     OpBuilder &builder, OperationState &result, Value target,
     ArrayRef<OpFoldResult> mixedTileSizes, transform::TileSizesSpec,
-    ArrayAttr mapping) {
+    ArrayAttr mappingAttr) {
   SmallVector<int64_t> staticTileSizes;
   SmallVector<Value> dynamicTileSizes;
   dispatchIndexOpFoldResults(mixedTileSizes, dynamicTileSizes, staticTileSizes,
                              ShapedType::kDynamicSize);
   // 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.
   MLIRContext *ctx = builder.getContext();
   auto operationType = pdl::OperationType::get(ctx);
   auto staticTileSizesAttr = builder.getI64ArrayAttr(staticTileSizes);
-  build(builder, result, TypeRange{operationType, operationType}, target,
-        /*numThreads=*/ValueRange{}, dynamicTileSizes,
-        /*staticNumThreads=*/ArrayAttr(), staticTileSizesAttr, mapping);
+  build(builder, result,
+        /*resultTypes=*/TypeRange{operationType, operationType},
+        /*target=*/target,
+        /*numThreads=*/ValueRange{},
+        /*tileSizes=*/dynamicTileSizes,
+        /*staticNumThreads=*/builder.getI64ArrayAttr({}),
+        /*staticTileSizes=*/staticTileSizesAttr,
+        /*mapping=*/mappingAttr);
+}
+
+void transform::TileToForeachThreadOp::build(OpBuilder &builder,
+                                             OperationState &result,
+                                             Value target,
+                                             Value packedTileSizesHandle,
+                                             transform::TileSizesSpec,
+                                             ArrayAttr mappingAttr) {
+  // 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.
+  MLIRContext *ctx = builder.getContext();
+  auto operationType = pdl::OperationType::get(ctx);
+  build(builder, result,
+        /*resultTypes=*/TypeRange{operationType, operationType},
+        /*target=*/target,
+        /*numThreads=*/ValueRange{},
+        /*tileSizes=*/ValueRange{packedTileSizesHandle},
+        /*staticNumThreads=*/builder.getI64ArrayAttr({}),
+        /*staticTileSizes=*/
+        builder.getI64ArrayAttr(ShapedType::kDynamicSize),
+        /*mapping=*/mappingAttr,
+        /*packedSizes=*/builder.getUnitAttr());
 }
 
 void transform::TileToForeachThreadOp::build(OpBuilder &builder,
                                              OperationState &result,
                                              Value target,
                                              ArrayRef<int64_t> staticNumThreads,
                                              transform::NumThreadsSpec,
-                                             ArrayAttr mapping) {
+                                             ArrayAttr mappingAttr) {
   return build(builder, result, target,
                getAsOpFoldResult(builder.getI64ArrayAttr(staticNumThreads)),
-               NumThreadsSpec(), mapping);
+               NumThreadsSpec(), mappingAttr);
 }
 
 void transform::TileToForeachThreadOp::build(
     OpBuilder &builder, OperationState &result, Value target,
     ArrayRef<OpFoldResult> mixedNumThreads, transform::NumThreadsSpec,
-    ArrayAttr mapping) {
+    ArrayAttr mappingAttr) {
   SmallVector<int64_t> staticNumThreads;
   SmallVector<Value> dynamicNumThreads;
   dispatchIndexOpFoldResults(mixedNumThreads, dynamicNumThreads,
                              staticNumThreads, ShapedType::kDynamicSize);
   // 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.
   MLIRContext *ctx = builder.getContext();
   auto operationType = pdl::OperationType::get(ctx);
   auto staticNumThreadsAttr = builder.getI64ArrayAttr(staticNumThreads);
-  build(builder, result, TypeRange{operationType, operationType}, target,
-        dynamicNumThreads, /*tileSizes=*/ValueRange{}, staticNumThreadsAttr,
-        /*staticTileSizes=*/ArrayAttr(), mapping);
+  build(builder, result,
+        /*resultTypes=*/TypeRange{operationType, operationType},
+        /*target=*/target,
+        /*numThreads=*/dynamicNumThreads,
+        /*tileSizes=*/ValueRange{},
+        /*staticNumThreads=*/staticNumThreadsAttr,
+        /*staticTileSizes=*/builder.getI64ArrayAttr({}),
+        /*mapping=*/mappingAttr);
 }
 
-DiagnosedSilenceableFailure transform::tileToForeachThreadOpImpl(
-    RewriterBase &rewriter, transform::TransformState &state,
-    TransformOpInterface transformOp, ArrayRef<Operation *> targets,
-    ArrayRef<OpFoldResult> mixedNumThreads,
-    ArrayRef<OpFoldResult> mixedTileSizes, Optional<ArrayAttr> mapping,
-    SmallVector<Operation *> &tileOps, SmallVector<Operation *> &tiledOps) {
-  if (targets.empty())
-    return DiagnosedSilenceableFailure(success());
-
-  // Given a list of OpFoldResults that are either index attrs or op handles,
-  // return a list of OpFoldResults where all op handles are replaced with the
-  // first (and only) OpResult of that payload op. (There must be exactly one
-  // mapped payload op and it must have exactly one index result.)
-  auto getOpResultsOrIndexAttrs =
-      [&](SmallVector<OpFoldResult> &result,
-          ArrayRef<OpFoldResult> opHandlesOrIndexAttrs) {
-        for (OpFoldResult ofr : opHandlesOrIndexAttrs) {
-          if (ofr.is<Attribute>()) {
-            result.push_back(ofr);
-            continue;
-          }
-          ArrayRef<Operation *> dynamicNumThreads =
-              state.getPayloadOps(ofr.get<Value>());
-          if (dynamicNumThreads.size() != 1) {
-            DiagnosedSilenceableFailure diag =
-                transformOp.emitSilenceableError()
-                << "handle must be mapped to exactly 1 payload op";
-            diag.attachNote(ofr.get<Value>().getLoc())
-                << "mapped to " << dynamicNumThreads.size() << " ops";
-            return diag;
-          }
-          Operation *op = dynamicNumThreads[0];
-          if (op->getNumResults() != 1 ||
-              !op->getResult(0).getType().isIndex()) {
-            DiagnosedSilenceableFailure diag =
-                transformOp.emitSilenceableError()
-                << "payload op must have exactly 1 index result";
-            diag.attachNote(op->getLoc())
-                << "has " << op->getNumResults() << " results";
-            return diag;
-          }
-          result.push_back(op->getResult(0));
-        }
+void transform::TileToForeachThreadOp::build(OpBuilder &builder,
+                                             OperationState &result,
+                                             Value target,
+                                             Value packedNumThreadsHandle,
+                                             transform::NumThreadsSpec,
+                                             ArrayAttr mappingAttr) {
+  // 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.
+  MLIRContext *ctx = builder.getContext();
+  auto operationType = pdl::OperationType::get(ctx);
+  build(builder, result,
+        /*resultTypes=*/TypeRange{operationType, operationType},
+        /*target=*/target,
+        /*numThreads=*/ValueRange{packedNumThreadsHandle},
+        /*tileSizes=*/ValueRange{},
+        /*staticNumThreads=*/
+        builder.getI64ArrayAttr(ShapedType::kDynamicSize),
+        /*staticTileSizes=*/builder.getI64ArrayAttr({}),
+        /*mapping=*/mappingAttr,
+        /*packedSizes=*/builder.getUnitAttr());
+}
+
+// Given a list of OpFoldResults that are either index attrs or op
+// handles, return a list of OpFoldResults where all op handles are
+// replaced with the first (and only) OpResult of that payload op. (There
+// must be exactly one mapped payload op and it must have exactly one
+// index result.)
+static DiagnosedSilenceableFailure unpackPDLOperations(
+    transform::TransformState &state, TransformOpInterface transformOp,
+    SmallVector<OpFoldResult> &result, ArrayRef<OpFoldResult> ofrs) {
+  for (OpFoldResult ofr : ofrs) {
+    // Don't try to unpack non-PDL operation.
+    if (ofr.is<Attribute>() ||
+        !ofr.get<Value>().getType().isa<pdl::OperationType>()) {
+    for (Operation *op : dynamicNumThreads) {
+      if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
+  }
 
-        return DiagnosedSilenceableFailure(success());
-      };
 
+DiagnosedSilenceableFailure transform::tileToForeachThreadOpImpl(
+    RewriterBase &rewriter, transform::TransformState &state,
+    TransformOpInterface transformOp, ArrayRef<Operation *> targets,
+    ArrayRef<OpFoldResult> mixedNumThreads,
+    ArrayRef<OpFoldResult> mixedTileSizes, Optional<ArrayAttr> mapping,
+    SmallVector<Operation *> &tileOps, SmallVector<Operation *> &tiledOps) {
+  if (targets.empty())
+    return DiagnosedSilenceableFailure(success());
+
   // getMixedNumThreads are OpFoldResults[index attributes or PDL operation].
   // Convert to OpFoldResults[index attributes or payload op].
   SmallVector<OpFoldResult> numThreads;
   DiagnosedSilenceableFailure status =
-      getOpResultsOrIndexAttrs(numThreads, mixedNumThreads);
+      unpackPDLOperations(state, transformOp, numThreads, mixedNumThreads);
   if (!status.succeeded())
     return status;
 
   // getMixedTileSizes are OpFoldResults[index attributes or PDL operation].
   // Convert to OpFoldResults[index attributes or payload op].
   SmallVector<OpFoldResult> tileSizes;
-  status = getOpResultsOrIndexAttrs(tileSizes, mixedTileSizes);
+  status = unpackPDLOperations(state, transformOp, tileSizes, mixedTileSizes);
   if (!status.succeeded())
     return status;
 
   // Transform all targets one by one.
   for (Operation *target : targets) {
     auto tilableOp = dyn_cast<TilingInterface>(target);
     if (!tilableOp) {
       DiagnosedSilenceableFailure diag =
DiagnosedSilenceableFailure transform::TileToForeachThreadOp::apply(
   SmallVector<Operation *> tileOps;
   SmallVector<Operation *> tiledOps;
 
   DiagnosedSilenceableFailure diag = tileToForeachThreadOpImpl(
       rewriter, state, cast<TransformOpInterface>(getOperation()), targets,
       getMixedNumThreads(), getMixedTileSizes(), getMapping(), tileOps,
       tiledOps);
 
-  if (!diag.succeeded())
+  if (!diag.succeeded()) {
+    transformResults.set(getForeachThreadOp().cast<OpResult>(), {});
+    transformResults.set(getTiledOp().cast<OpResult>(), {});
     return diag;
+  }
 
   transformResults.set(getForeachThreadOp().cast<OpResult>(), tileOps);
   transformResults.set(getTiledOp().cast<OpResult>(), tiledOps);
 
   return DiagnosedSilenceableFailure(success());
 }
 
 void transform::TileToForeachThreadOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   onlyReadsHandle(getTileSizes(), effects);
   onlyReadsHandle(getNumThreads(), effects);
   producesHandle(getResults(), effects);
 }
 
 SmallVector<OpFoldResult> TileToForeachThreadOp::getMixedNumThreads() {
   return getMixedSizes(getStaticNumThreads(), getNumThreads());
 }
 
 SmallVector<OpFoldResult> TileToForeachThreadOp::getMixedTileSizes() {
   return getMixedSizes(getStaticTileSizes(), getTileSizes());
 }
 
 LogicalResult TileToForeachThreadOp::verify() {
-  if (getMixedNumThreads().empty() == getMixedTileSizes().empty())
+  if (getMixedNumThreads().empty() == getMixedTileSizes().empty()) {
     return emitOpError("either num_threads or tile_sizes must be specified");
+  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TileToScfForOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir

//  CHECK-SAME:   %[[C:[0-9a-z]+]]: tensor<?x?xf32>
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
     %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20] (mapping = [ #gpu.thread<y>, #gpu.thread<x> ] )
   }
 }
 
 // -----
 
+// In this test case, matmul dims and tile size are dynamic.
+
+// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
+// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
+// CHECK-DAG: #[[$map4:.+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+
+// CHECK-LABEL: matmul_tile_size_dynamic_dynamic(
+//  CHECK-SAME:   %[[A:[0-9a-z]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[B:[0-9a-z]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[C:[0-9a-z]+]]: tensor<?x?xf32>
+func.func @matmul_tile_size_dynamic_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  //  CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
+  //  CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
+  //  CHECK-DAG: %[[tile_size_1:.*]] = "test.dummy"()
+  //  CHECK-DAG: %[[tile_size_2:.*]] = "test.dummy"()
+  //  CHECK-DAG: %[[M:.+]] = tensor.dim %[[A]], %[[c0]] :
+  //  CHECK-DAG: %[[N:.+]] = tensor.dim %[[B]], %c1 :
+  //  CHECK-DAG: %[[NT0:.+]] = affine.apply #[[$map0]]()[%[[M]], %[[tile_size_1]]]
+  //  CHECK-DAG: %[[NT1:.+]] = affine.apply #[[$map0]]()[%[[N]], %[[tile_size_2]]]
+  //      CHECK: scf.foreach_thread (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]]) shared_outs(%[[C_BLK:.*]] = %[[C]])
+  //      CHECK:   tensor.extract_slice %[[A]]
+  //      CHECK:   tensor.extract_slice %[[B]]
+  //      CHECK:   tensor.extract_slice %[[C_BLK]]
+  //      CHECK:   linalg.matmul
+  //      CHECK:   scf.foreach_thread.perform_concurrently
+  // CHECK-NEXT:    tensor.parallel_insert_slice
+  %tile_size_1 = "test.dummy"() : () -> (index)
+  %tile_size_2 = "test.dummy"() : () -> (index)
+  %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+                    outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
+  return %0 : tensor<?x?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
+  %sz = transform.structured.match ops{["test.dummy"]} in %arg1
+  %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [%sz]
+}
+
+// -----
+
 // Tests that dimension 0 can eliminate affine.min/max, dimension 1 cannot.
 
 // CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * -15 + 300, 15)>
 // CHECK-DAG: #[[$map1:.+]] = affine_map<(d0) -> (0, d0)>
 // CHECK-DAG: #[[$map2:.+]] = affine_map<(d0) -> (d0 * 10)>
 // CHECK-DAG: #[[$map3:.+]] = affine_map<(d0) -> (d0 * 15)>
 
 // CHECK-LABEL: matmul_static(