# RUN: %PYTHON %s | FileCheck %s

from mlir.ir import *
from mlir.dialects import builtin
from mlir.dialects import linalg
from mlir.dialects import std

def run(f):
  print("\nTEST:", f.__name__)
  f()
  return f


# CHECK-LABEL: TEST: testInitTensor
@run
def testInitTensor():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):
      # CHECK-LABEL: func @static_sizes
      # CHECK: %0 = linalg.init_tensor [3, 4] : tensor<3x4xf32>
      @builtin.FuncOp.from_py_func()
      def static_sizes():
        return linalg.InitTensorOp([3, 4], f32)

      # CHECK-LABEL: func @dynamic_sizes
      # CHECK: %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
      @builtin.FuncOp.from_py_func(IndexType.get(), IndexType.get())
      def dynamic_sizes(d0, d1):
        return linalg.InitTensorOp([d0, d1], f32)

      # CHECK-LABEL: func @zero_d
      # CHECK: %0 = linalg.init_tensor [] : tensor<f32>
      @builtin.FuncOp.from_py_func()
      def zero_d():
        return linalg.InitTensorOp([], f32)

  print(module)

# CHECK-LABEL: TEST: testInitTensorStaticSizesAttribute
@run
def testInitTensorStaticSizesAttribute():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):
      op = linalg.InitTensorOp([3, 4], f32)
      # CHECK: [3, 4]
      print(op.attributes['static_sizes'])

# CHECK-LABEL: TEST: testFill
@run
def testFill():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):
      # CHECK-LABEL: func @fill_tensor
      #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: tensor<12x?xf32>
      #  CHECK-NEXT: %[[CST:.*]] = constant 0.0{{.*}} : f32
      #  CHECK-NEXT: %[[RES:.*]] = linalg.fill(%[[CST]], %[[OUT]]) : f32, tensor<12x?xf32> -> tensor<12x?xf32>
      #  CHECK-NEXT: return %[[RES]] : tensor<12x?xf32>
      @builtin.FuncOp.from_py_func(
          RankedTensorType.get((12, -1), f32))
      def fill_tensor(out):
        zero = std.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
        # TODO: FillOp.result is None. When len(results) == 1 we expect it to
        # be results[0] as per _linalg_ops_gen.py. This seems like an
        # orthogonal bug in the generator of _linalg_ops_gen.py.
        return linalg.FillOp(output=out, value=zero).results[0]

      # CHECK-LABEL: func @fill_buffer
      #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: memref<12x?xf32>
      #  CHECK-NEXT: %[[CST:.*]] = constant 0.0{{.*}} : f32
      #  CHECK-NEXT: linalg.fill(%[[CST]], %[[OUT]]) : f32, memref<12x?xf32>
      #  CHECK-NEXT: return
      @builtin.FuncOp.from_py_func(
          MemRefType.get((12, -1), f32))
      def fill_buffer(out):
        zero = std.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
        linalg.FillOp(output=out, value=zero)

  print(module)


# CHECK-LABEL: TEST: testStructuredOpOnTensors
@run
def testStructuredOpOnTensors():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    tensor_type = RankedTensorType.get((2, 3, 4), f32)
    with InsertionPoint(module.body):
      func = builtin.FuncOp(name="matmul_test",
                            type=FunctionType.get(
                                inputs=[tensor_type, tensor_type],
                                results=[tensor_type]))
      with InsertionPoint(func.add_entry_block()):
        lhs, rhs = func.entry_block.arguments
        result = linalg.MatmulOp([lhs, rhs], results=[tensor_type]).result
        std.ReturnOp([result])

  # CHECK: %[[R:.*]] = linalg.matmul ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
  print(module)


# CHECK-LABEL: TEST: testStructuredOpOnBuffers
@run
def testStructuredOpOnBuffers():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    memref_type = MemRefType.get((2, 3, 4), f32)
    with InsertionPoint(module.body):
      func = builtin.FuncOp(name="matmul_test",
                            type=FunctionType.get(
                                inputs=[memref_type, memref_type, memref_type],
                                results=[]))
      with InsertionPoint(func.add_entry_block()):
        lhs, rhs, result = func.entry_block.arguments
        # TODO: prperly hook up the region.
        linalg.MatmulOp([lhs, rhs], outputs=[result])
        std.ReturnOp([])

  # CHECK: linalg.matmul ins(%arg0, %arg1 : memref<2x3x4xf32>, memref<2x3x4xf32>) outs(%arg2 : memref<2x3x4xf32>)
  print(module)

# CHECK-LABEL: TEST: testNamedStructuredOpCustomForm
@run
def testNamedStructuredOpCustomForm():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):
      @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                   RankedTensorType.get((16, 8), f32))
      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])

  print(module)

# CHECK-LABEL: TEST: testNamedStructuredOpGenericForm
@run
def testNamedStructuredOpGenericForm():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):
      @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                   RankedTensorType.get((16, 8), f32))
      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])

  module.operation.print(print_generic_op_form=True)

# CHECK-LABEL: TEST: testNamedStructuredAsGenericOp
@run
def testNamedStructuredAsGenericOp():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):
      @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                   RankedTensorType.get((16, 8), f32))
      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)

  print(module)


# CHECK-LABEL: TEST: testOpResultFromOtherOp
@run
def testOpResultFromOtherOp():
  with Context(), Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):

      @builtin.FuncOp.from_py_func(
          RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
                                                                   f32))
      def pass_an_op_directly(arg0, arg1):
        one = std.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)

  print(module)