def generic_form(lhs, rhs): init_result = linalg.InitTensorOp([4, 8], f32) # CHECK: linalg.generic return linalg.matmul(lhs, rhs, outs=[init_result.result], emit_generic=True)
def named_form(lhs, rhs): init_result = linalg.InitTensorOp([4, 8], f32) # CHECK: "linalg.matmul"(%{{.*}}) # CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32): # CHECK-NEXT: std.mulf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: std.addf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: linalg.yield{{.*}} (f32) -> () # CHECK-NEXT: {linalg.memoized_indexing_maps{{.*}}operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} : # CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> return linalg.matmul(lhs, rhs, outs=[init_result.result])
def named_form(lhs, rhs): init_result = linalg.InitTensorOp([4, 8], f32) # First check the named form with custom format # CHECK: linalg.matmul # CHECK-NOT: linalg.memoized_indexing_maps # CHECK-SAME: ins(%{{.*}} : tensor<4x16xf32>, tensor<16x8xf32>) # CHECK-SAME: outs(%{{.*}} : tensor<4x8xf32>) # CHECK-SAME: -> tensor<4x8xf32> # CHECK-NEXT: return return linalg.matmul(lhs, rhs, outs=[init_result.result])
def named_form(lhs, rhs): init_result = linalg.InitTensorOp([4, 8], f32) # CHECK: "linalg.matmul"(%{{.*}}) # CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32): # CHECK-NEXT: arith.mulf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: arith.addf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: linalg.yield{{.*}} (f32) -> () # CHECK-NEXT: cast = #linalg.type_fn<cast_signed> # CHECK-SAME: operand_segment_sizes = dense<[2, 1]> : vector<2xi32> # CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> return linalg.matmul(lhs, rhs, outs=[init_result.result])
def pass_an_op_directly(arg0, arg1): one = arith.ConstantOp(F32Type.get(), 1.0) # CHECK: %[[LHS:.*]] = linalg.fill lhs = linalg.FillOp(arg0, one) # CHECK: %[[RHS:.*]] = linalg.fill rhs = linalg.FillOp(arg1, one) # CHECK: %[[INIT:.*]] = linalg.init_tensor init = linalg.InitTensorOp([4, 8], f32) # CHECK: linalg.matmul # CHECK: ins(%[[LHS]], %[[RHS]] # CHECK: outs(%[[INIT]] return linalg.matmul(lhs, rhs, outs=init)
def matmul_on_buffers(lhs, rhs, out): linalg.matmul(lhs, rhs, outs=[out])
def matmul_on_buffers(lhs, rhs, out): linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)
def matmul_unsigned_on_buffers(lhs, rhs, out): linalg.matmul(lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned, emit_generic=True)
def matmul_unsigned_on_buffers(lhs, rhs, out): linalg.matmul(lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned)