bjacob

So with this:

  LogicalResult matchAndRewrite(linalg::MatmulOp op,
                                PatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<meh::MatmulOp>(op, op.getOperand(0),
                                               op.getOperand(1));
    return success();
  }

I get this build error:

bjacob

// meh::MatmulOp takes two Tensors as arguments lhs() and rhs(), and returns a Tensor, dst().

// The below is code that was originally written against memrefs, and doesn't work anymore now what meh::MatmulOp deals with tensors.

// The question is how to fix it to work with tensors.

class MehMatmulToSCFPattern : public OpRewritePattern<meh::MatmulOp> {
 public:
  using OpRewritePattern<meh::MatmulOp>::OpRewritePattern;

bjacob

// I'm trying to generate a simple loop nest performing matrix multiplication,
// while staying at the Tensor level.

// Everything worked well while this code was working at the memref level.
// But now I don't know how to end this rewrite pass:

// I can't eraseOp because it says that the old op's result is still used.

// If I don't do anything at the end, my pass has no effect, presumably
// because I haven't actually replaced the old op by the new loop nest.

bjacob

<stdin>:4:10: error: operand #3 does not dominate this use
    %0 = linalg.matmul ins(%lhs, %rhs: tensor<?x?xf32>, tensor<?x?xf32>)
         ^
<stdin>:4:10: note: see current operation: %0 = "scf.for"(%c0, %c-1, %c1, %0) ( {
^bb0(%arg3: index, %arg4: tensor<?x?xf32>):  // no predecessors
  %1 = "scf.for"(%c0, %c-1, %c1, %0) ( {
  ^bb0(%arg5: index, %arg6: tensor<?x?xf32>):  // no predecessors
    %2 = "scf.for"(%c0, %c-1, %c1, %0) ( {
    ^bb0(%arg7: index, %arg8: tensor<?x?xf32>):  // no predecessors
      %3 = "std.subtensor"(%arg0, %arg3, %arg7, %c1, %c1, %c1, %c1) {operand_segment_sizes = dense<[1, 2, 2, 2]> : vector<4xi32>, static_offsets = [-9223372036854775808, -9223372036854775808], static_sizes = [-1, -1], static_strides = [-9223372036854775808, -9223372036854775808]} : (tensor<?x?xf32>, index, index, index, index, index, index) -> tensor<?x?xf32>

bjacob

You can load the mnist.tracy trace directly in the Tracy UI.

To reproduce:

• download mnist.mlir

bjacob

Here is a matmul with two ops, producer_lhs and producer_rhs, fused into it. The producers have a cost. They could be just reading constant data (e.g. weights of a conv op) or they could be more expensive math (e.g. math-function activation function of preceding layer). Either way, they have non-negligible cost (even reading constant data has the cost of memory accesses).

for (int i = 0; i < M; i++) {
  for (int j = 0; j < N; j++) {
    for (int k = 0; k < K; k++) {
 result[i, j] += producer_lhs(i, k) * producer_rhs(k, j);

bjacob

Starting point: In IREE/Linalg,

1. There ultimately won't be any packing ops (because anything resembling a packing op will get fused into the producer).
2. But there will still be a packed layout (because the repeated N^3 accesses to N^2 data mean that it is an optimization to materialize the ultimate matmul lhs/rhs operands as buffers, and efficiency considerations mean that the layout of these buffers will have some block structure).

So a priori, at some stage of lowering just before codegen, our matmuls have as lhs/rhs inputs some materialized buffers in some possibly nontrivial layouts.

I started thinking about how much we can restrict layouts. If you look at classic GEMM papers like the BLIS papers, it can scare you because they make it sound like you need N nested levels of blocks, and some of that may need to reflect the CPU cache hierarchy.

And so in particular that makes it sound like the introduction of packed layouts is adding new degrees of freedom, new dimensions to the search space in th

bjacob

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index 5752af9bea9a..b56b6ffa9c86 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -305,4 +305,62 @@ structured_op: !LinalgStructuredOpConfig
                 operands:
                 - !ScalarExpression
                   scalar_arg: B
-
+--- !LinalgOpConfig

bjacob

yaml:

--- !LinalgOpConfig
metadata: !LinalgOpMetadata
  name: mmt_kernel
  cpp_op_name: MmtKernelOp
  doc: |-
    WRITE ME
  implements:

bjacob

x.mlir is generated by mlir-proto-opt -convert-scf-to-std.

Next, I try:

$ cat /tmp/x.mlir | blaze-bin/third_party/iree/experimental/runners/mlir-proto-opt -convert-std-to-llvm 
<stdin>:59:11: error: 'std.tensor_load' op requires the same shape for all operands and results
    %14 = tensor_load %1 : memref<2x4x8x2xf32>
          ^
<stdin>:59:11: note: see current operation: %85 = "std.tensor_load"(%55) : (!llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>) -> tensor<2x4x8x2xf32>