README.md

Command line:

blaze-bin/third_party/iree/experimental/runners/mlir-proto-opt -linalg-comprehensive-bufferize-inplace /tmp/a.mlir

Output:

return val does not fold: %0 = tensor.generate %arg0, %arg1, %arg2, %arg3  {
^bb0(%arg4: index, %arg5: index, %arg6: index, %arg7: index):  // no predecessors
  %1 = index_cast %arg4 : index to i32
  %2 = sitofp %1 : i32 to f32
  tensor.yield %2 : f32
} : tensor<?x?x?x?xf32>
/tmp/a.mlir:26:10: error: 'std.tensor_load' op requires the same shape for all operands and results
  %lhs = call @generate_pseudorandom_4d_f32 (%M, %M0, %K, %K0) : (index, index, index, index) -> tensor<?x?x?x?xf32>
         ^
/tmp/a.mlir:26:10: note: see current operation: %1 = "std.tensor_load"(%0) : (tensor<?x?x?x?xf32>) -> none
/tmp/a.mlir:26:10: error: 'std.tensor_load' op requires the same shape for all operands and results
  %lhs = call @generate_pseudorandom_4d_f32 (%M, %M0, %K, %K0) : (index, index, index, index) -> tensor<?x?x?x?xf32>
         ^
/tmp/a.mlir:26:10: note: see current operation: %1 = "std.tensor_load"(%0) : (tensor<?x?x?x?xf32>) -> none

a.mlir

!elem_type_lhs = type f32
!elem_type_rhs = type f32
!elem_type_accum = type f32
!tensor_type_lhs = type tensor<2x4x8x2x!elem_type_lhs>
!tensor_type_rhs = type tensor<3x4x4x2x!elem_type_rhs>
!tensor_type_accum = type tensor<2x3x8x4x!elem_type_accum>
!vector_type_accum = type vector<2x3x8x4x!elem_type_accum>

func @generate_pseudorandom_4d_f32(%s0 : index, %s1 : index, %s2 : index, %s3 : index) -> tensor<?x?x?x?xf32> {
  %result = tensor.generate %s0, %s1, %s2, %s3 {
    ^bb0(%i0 : index, %i1 : index, %i2 : index, %i3 : index):
    %i_i32 = index_cast %i0 : index to i32
    %i_f32 = sitofp %i_i32 : i32 to f32
    tensor.yield %i_f32 : f32
  } : tensor<?x?x?x?xf32>
  return %result : tensor<?x?x?x?xf32>
}

func @main() {
  %M = constant 2 : index
  %M0 = constant 8 : index
  %K = constant 4 : index
  %K0 = constant 2 : index
  %MM0 = std.muli %M, %M0 : index
  %KK0 = std.muli %K, %K0 : index
  %lhs = call @generate_pseudorandom_4d_f32 (%M, %M0, %K, %K0) : (index, index, index, index) -> tensor<?x?x?x?xf32>
  %rhs = tensor.generate {
    ^bb0(%n : index, %k : index, %n0 : index, %k0 : index):
    %i = affine.apply affine_map<(n, k, n0, k0) -> (n + k + n0 + k0)> (%n, %k, %n0, %k0)
    %i_i32 = index_cast %i : index to i32
    %i_f32 = sitofp %i_i32 : i32 to f32
    tensor.yield %i_f32 : f32
  } : !tensor_type_rhs
  %accum = tensor.generate {
    ^bb0(%m : index, %n : index, %m0 : index, %n0 : index):
    %i = affine.apply affine_map<(m, n, m0, n0) -> (m + n + m0 + n0)> (%m, %n, %m0, %n0)
    %i_i32 = index_cast %i : index to i32
    %i_f32 = sitofp %i_i32 : i32 to f32
    tensor.yield %i_f32 : f32
  } : !tensor_type_accum

  %result = linalg.mmt_4d_kernel ins(%lhs, %rhs : tensor<?x?x?x?xf32>, !tensor_type_rhs)
    outs(%accum: !tensor_type_accum) -> !tensor_type_accum

  %c0 = constant 0: index
  %v0 = constant 0.0 : !elem_type_accum
  %result_vector = vector.transfer_read %result[%c0, %c0, %c0, %c0], %v0 : !tensor_type_accum, !vector_type_accum
  vector.print %result_vector: !vector_type_accum

  return
}