# Copyright 2019 The MLIR Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# RUN: %p/test_edsc %s | FileCheck %s
"""Python2 and 3 test for the MLIR EDSC Python bindings"""

import google_mlir.bindings.python.pybind as E
import inspect

# Prints `str` prefixed by the current test function name so we can use it in
# Filecheck label directives.
# This is achieved by inspecting the stack and getting the parent name.
def printWithCurrentFunctionName(str):
  print(inspect.stack()[1][3])
  print(str)

class EdscTest:

  def setUp(self):
    self.module = E.MLIRModule()
    self.boolType = self.module.make_type("i1")
    self.i32Type = self.module.make_type("i32")
    self.f32Type = self.module.make_type("f32")
    self.indexType = self.module.make_index_type()

  def testBlockArguments(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      E.constant_index(42)
      with E.BlockContext([self.f32Type, self.f32Type]) as b:
        b.arg(0) + b.arg(1)
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBlockArguments
    #       CHECK: %{{.*}} = constant 42 : index
    #       CHECK: ^bb{{.*}}(%{{.*}}: f32, %{{.*}}: f32):
    #       CHECK:   %{{.*}} = addf %{{.*}}, %{{.*}} : f32

  def testBlockContext(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      cst = E.constant_index(42)
      with E.BlockContext():
        cst + cst
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBlockContext
    #       CHECK: %{{.*}} = constant 42 : index
    #       CHECK: ^bb
    #       CHECK: %{{.*}} = "affine.apply"() {map = () -> (84)} : () -> index

  def testBlockContextAppend(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      E.constant_index(41)
      with E.BlockContext() as b:
        blk = b  # save block handle for later
        E.constant_index(0)
      E.constant_index(42)
      with E.BlockContext(E.appendTo(blk)):
        E.constant_index(1)
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBlockContextAppend
    #       CHECK: %{{.*}} = constant 41 : index
    #       CHECK: %{{.*}} = constant 42 : index
    #       CHECK: ^bb
    #       CHECK: %{{.*}} = constant 0 : index
    #       CHECK: %{{.*}} = constant 1 : index

  def testBlockContextStandalone(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      blk1 = E.BlockContext()
      blk2 = E.BlockContext()
      with blk1:
        E.constant_index(0)
      with blk2:
        E.constant_index(56)
        E.constant_index(57)
      E.constant_index(41)
      with blk1:
        E.constant_index(1)
      E.constant_index(42)
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBlockContextStandalone
    #       CHECK: %{{.*}} = constant 41 : index
    #       CHECK: %{{.*}} = constant 42 : index
    #       CHECK: ^bb
    #       CHECK: %{{.*}} = constant 0 : index
    #       CHECK: %{{.*}} = constant 1 : index
    #       CHECK: ^bb
    #       CHECK: %{{.*}} = constant 56 : index
    #       CHECK: %{{.*}} = constant 57 : index

  def testBooleanOps(self):
    self.setUp()
    with self.module.function_context(
        "booleans", [self.boolType for _ in range(4)], []) as fun:
      i, j, k, l = (fun.arg(x) for x in range(4))
      stmt1 = (i < j) & (j >= k)
      stmt2 = ~(stmt1 | (k == l))
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBooleanOps
    #       CHECK: %{{.*}} = cmpi "slt", %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = cmpi "sge", %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = muli %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = cmpi "eq", %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = constant 1 : i1
    #       CHECK: %{{.*}} = subi %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = constant 1 : i1
    #       CHECK: %{{.*}} = subi %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = muli %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = constant 1 : i1
    #       CHECK: %{{.*}} = subi %{{.*}}, %{{.*}} : i1
    #       CHECK: %{{.*}} = constant 1 : i1
    #       CHECK: %{{.*}} = subi %{{.*}}, %{{.*}} : i1

  def testBr(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      with E.BlockContext() as b:
        blk = b
        E.ret()
      E.br(blk)
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBr
    #       CHECK:   br ^bb
    #       CHECK: ^bb
    #       CHECK:   return

  def testBrArgs(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      # Create an infinite loop.
      with E.BlockContext([self.indexType, self.indexType]) as b:
        E.br(b, [b.arg(1), b.arg(0)])
      E.br(b, [E.constant_index(0), E.constant_index(1)])
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBrArgs
    #       CHECK:   %{{.*}} = constant 0 : index
    #       CHECK:   %{{.*}} = constant 1 : index
    #       CHECK:   br ^bb{{.*}}(%{{.*}}, %{{.*}} : index, index)
    #       CHECK: ^bb{{.*}}(%{{.*}}: index, %{{.*}}: index):
    #       CHECK:   br ^bb{{.*}}(%{{.*}}, %{{.*}} : index, index)

  def testBrDeclaration(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      blk = E.BlockContext()
      E.br(blk.handle())
      with blk:
        E.ret()
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testBrDeclaration
    #       CHECK:   br ^bb
    #       CHECK: ^bb
    #       CHECK:   return

  def testCallOp(self):
    self.setUp()
    callee = self.module.declare_function("sqrtf", [self.f32Type],
                                          [self.f32Type])
    with self.module.function_context("call", [self.f32Type], []) as fun:
      funCst = E.constant_function(callee)
      funCst([fun.arg(0)]) + E.constant_float(42., self.f32Type)
      printWithCurrentFunctionName(str(self.module))
    # CHECK-LABEL: testCallOp
    #       CHECK: func @sqrtf(f32) -> f32
    #       CHECK:   %{{.*}} = constant @sqrtf : (f32) -> f32
    #       CHECK:   %{{.*}} = call_indirect %{{.*}}(%{{.*}}) : (f32) -> f32

  def testCondBr(self):
    self.setUp()
    with self.module.function_context("foo", [self.boolType], []) as fun:
      with E.BlockContext() as blk1:
        E.ret([])
      with E.BlockContext([self.indexType]) as blk2:
        E.ret([])
      cst = E.constant_index(0)
      E.cond_br(fun.arg(0), blk1, [], blk2, [cst])
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testCondBr
    #       CHECK:   cond_br %{{.*}}, ^bb{{.*}}, ^bb{{.*}}(%{{.*}} : index)

  def testConstants(self):
    self.setUp()
    with self.module.function_context("constants", [], []) as fun:
      E.constant_float(1.23, self.module.make_type("bf16"))
      E.constant_float(1.23, self.module.make_type("f16"))
      E.constant_float(1.23, self.module.make_type("f32"))
      E.constant_float(1.23, self.module.make_type("f64"))
      E.constant_int(1, 1)
      E.constant_int(123, 8)
      E.constant_int(123, 16)
      E.constant_int(123, 32)
      E.constant_int(123, 64)
      E.constant_index(123)
      E.constant_function(fun)
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testConstants
    #       CHECK:  constant 1.230000e+00 : bf16
    #       CHECK:  constant 1.230470e+00 : f16
    #       CHECK:  constant 1.230000e+00 : f32
    #       CHECK:  constant 1.230000e+00 : f64
    #       CHECK:  constant 1 : i1
    #       CHECK:  constant 123 : i8
    #       CHECK:  constant 123 : i16
    #       CHECK:  constant 123 : i32
    #       CHECK:  constant 123 : index
    #       CHECK:  constant @constants : () -> ()

  def testCustom(self):
    self.setUp()
    with self.module.function_context("custom", [self.indexType, self.f32Type],
                                      []) as fun:
      E.op("foo", [fun.arg(0)], [self.f32Type]) + fun.arg(1)
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testCustom
    #       CHECK: %{{.*}} = "foo"(%{{.*}}) : (index) -> f32
    #       CHECK:  %{{.*}} = addf %{{.*}}, %{{.*}} : f32

  # Create 'addi' using the generic Op interface.  We need an operation known
  # to the execution engine so that the engine can compile it.
  def testCustomOpCompilation(self):
    self.setUp()
    with self.module.function_context("adder", [self.i32Type], []) as f:
      c1 = E.op(
          "std.constant", [], [self.i32Type],
          value=self.module.integerAttr(self.i32Type, 42))
      E.op("std.addi", [c1, f.arg(0)], [self.i32Type])
      E.ret([])
    self.module.compile()
    printWithCurrentFunctionName(str(self.module.get_engine_address() == 0))
    # CHECK-LABEL: testCustomOpCompilation
    #       CHECK: False

  def testDivisions(self):
    self.setUp()
    with self.module.function_context(
        "division", [self.indexType, self.i32Type, self.i32Type], []) as fun:
      # indices only support floor division
      fun.arg(0) // E.constant_index(42)
      # regular values only support regular division
      fun.arg(1) / fun.arg(2)
      printWithCurrentFunctionName(str(self.module))
    # CHECK-LABEL: testDivisions
    #       CHECK:  floordiv 42
    #       CHECK:  divis %{{.*}}, %{{.*}} : i32

  def testFunctionArgs(self):
    self.setUp()
    with self.module.function_context("foo", [self.f32Type, self.f32Type],
                                      [self.indexType]) as fun:
      pass
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testFunctionArgs
    #       CHECK: func @foo(%{{.*}}: f32, %{{.*}}: f32) -> index

  def testFunctionContext(self):
    self.setUp()
    with self.module.function_context("foo", [], []):
      pass
      printWithCurrentFunctionName(self.module.get_function("foo"))
    # CHECK-LABEL: testFunctionContext
    #       CHECK: func @foo() {

  def testFunctionDeclaration(self):
    self.setUp()
    boolAttr = self.module.boolAttr(True)
    t = self.module.make_memref_type(self.f32Type, [10])
    t_llvm_noalias = t({"llvm.noalias": boolAttr})
    t_readonly = t({"readonly": boolAttr})
    f = self.module.declare_function("foo", [t, t_llvm_noalias, t_readonly], [])
    printWithCurrentFunctionName(str(self.module))
    # CHECK-LABEL: testFunctionDeclaration
    #       CHECK: func @foo(memref<10xf32>, memref<10xf32> {llvm.noalias = true}, memref<10xf32> {readonly = true})

  def testFunctionMultiple(self):
    self.setUp()
    with self.module.function_context("foo", [], []):
      pass
    with self.module.function_context("foo", [], []):
      E.constant_index(0)
    printWithCurrentFunctionName(str(self.module))
    # CHECK-LABEL: testFunctionMultiple
    #       CHECK: func @foo()
    #       CHECK: func @foo_0()
    #       CHECK: %{{.*}} = constant 0 : index

  def testIndexCast(self):
    self.setUp()
    with self.module.function_context("testIndexCast", [], []):
      index = E.constant_index(0)
      E.index_cast(index, self.i32Type)
    printWithCurrentFunctionName(str(self.module))
    # CHECK-LABEL: testIndexCast
    #       CHECK: index_cast %{{.*}} : index to i32

  def testIndexedValue(self):
    self.setUp()
    memrefType = self.module.make_memref_type(self.f32Type, [10, 42])
    with self.module.function_context("indexed", [memrefType],
                                      [memrefType]) as fun:
      A = E.IndexedValue(fun.arg(0))
      cst = E.constant_float(1., self.f32Type)
      with E.LoopNestContext(
          [E.constant_index(0), E.constant_index(0)],
          [E.constant_index(10), E.constant_index(42)], [1, 1]) as (i, j):
        A.store([i, j], A.load([i, j]) + cst)
      E.ret([fun.arg(0)])
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testIndexedValue
    #       CHECK: "affine.for"()
    #       CHECK: "affine.for"()
    #       CHECK: "affine.load"
    #  CHECK-SAME: memref<10x42xf32>
    #       CHECK:  %{{.*}} = addf %{{.*}}, %{{.*}} : f32
    #       CHECK:  "affine.store"
    #  CHECK-SAME:  memref<10x42xf32>
    #       CHECK: {lower_bound = () -> (0), step = 1 : index, upper_bound = () -> (42)}
    #       CHECK: {lower_bound = () -> (0), step = 1 : index, upper_bound = () -> (10)}

  def testLoopContext(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      lhs = E.constant_index(0)
      rhs = E.constant_index(42)
      with E.LoopContext(lhs, rhs, 1) as i:
        lhs + rhs + i
        with E.LoopContext(rhs, rhs + rhs, 2) as j:
          x = i + j
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testLoopContext
    #       CHECK: "affine.for"() (
    #       CHECK:   ^bb{{.*}}(%{{.*}}: index):
    #       CHECK: "affine.for"(%{{.*}}, %{{.*}}) (
    #       CHECK: ^bb{{.*}}(%{{.*}}: index):
    #       CHECK: "affine.apply"(%{{.*}}, %{{.*}}) {map = (d0, d1) -> (d0 + d1)} : (index, index) -> index
    #       CHECK: {lower_bound = (d0) -> (d0), step = 2 : index, upper_bound = (d0) -> (d0)} : (index, index) -> ()
    #       CHECK: {lower_bound = () -> (0), step = 1 : index, upper_bound = () -> (42)}

  def testLoopNestContext(self):
    self.setUp()
    with self.module.function_context("foo", [], []) as fun:
      lbs = [E.constant_index(i) for i in range(4)]
      ubs = [E.constant_index(10 * i + 5) for i in range(4)]
      with E.LoopNestContext(lbs, ubs, [1, 3, 5, 7]) as (i, j, k, l):
        i + j + k + l
    printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testLoopNestContext
    #       CHECK: "affine.for"() (
    #       CHECK: ^bb{{.*}}(%{{.*}}: index):
    #       CHECK: "affine.for"() (
    #       CHECK: ^bb{{.*}}(%{{.*}}: index):
    #       CHECK: "affine.for"() (
    #       CHECK: ^bb{{.*}}(%{{.*}}: index):
    #       CHECK: "affine.for"() (
    #       CHECK: ^bb{{.*}}(%{{.*}}: index):
    #       CHECK: %{{.*}} = "affine.apply"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {map = (d0, d1, d2, d3) -> (d0 + d1 + d2 + d3)} : (index, index, index, index) -> index

  def testMLIRBooleanCompilation(self):
    self.setUp()
    m = self.module.make_memref_type(self.boolType, [10])  # i1 tensor
    with self.module.function_context("mkbooltensor", [m, m], []) as f:
      input = E.IndexedValue(f.arg(0))
      output = E.IndexedValue(f.arg(1))
      zero = E.constant_index(0)
      ten = E.constant_index(10)
      with E.LoopNestContext([zero] * 3, [ten] * 3, [1] * 3) as (i, j, k):
        b1 = (i < j) & (j < k)
        b2 = ~b1
        b3 = b2 | (k < j)
        output.store([i], input.load([i]) & b3)
      E.ret([])
    self.module.compile()
    printWithCurrentFunctionName(str(self.module.get_engine_address() == 0))
    # CHECK-LABEL: testMLIRBooleanCompilation
    #       CHECK: False

  def testMLIRFunctionCreation(self):
    self.setUp()
    module = E.MLIRModule()
    t = module.make_type("f32")
    m = module.make_memref_type(t, [3, 4, -1, 5])
    printWithCurrentFunctionName(str(t))
    print(str(m))
    print(str(module.make_function("copy", [m, m], [])))
    print(str(module.make_function("sqrtf", [t], [t])))
    # CHECK-LABEL: testMLIRFunctionCreation
    #       CHECK:  f32
    #       CHECK:  memref<3x4x?x5xf32>
    #       CHECK: func @copy(%{{.*}}: memref<3x4x?x5xf32>, %{{.*}}: memref<3x4x?x5xf32>) {
    #       CHECK:  func @sqrtf(%{{.*}}: f32) -> f32

  def testMLIRScalarTypes(self):
    self.setUp()
    module = E.MLIRModule()
    printWithCurrentFunctionName(str(module.make_type("bf16")))
    print(str(module.make_type("f16")))
    print(str(module.make_type("f32")))
    print(str(module.make_type("f64")))
    print(str(module.make_type("i1")))
    print(str(module.make_type("i8")))
    print(str(module.make_type("i32")))
    print(str(module.make_type("i123")))
    print(str(module.make_type("index")))
    # CHECK-LABEL: testMLIRScalarTypes
    #       CHECK:  bf16
    #       CHECK:  f16
    #       CHECK:  f32
    #       CHECK:  f64
    #       CHECK:  i1
    #       CHECK:  i8
    #       CHECK:  i32
    #       CHECK:  i123
    #       CHECK:  index

  def testMatrixMultiply(self):
    self.setUp()
    memrefType = self.module.make_memref_type(self.f32Type, [32, 32])
    with self.module.function_context(
        "matmul", [memrefType, memrefType, memrefType], []) as fun:
      A = E.IndexedValue(fun.arg(0))
      B = E.IndexedValue(fun.arg(1))
      C = E.IndexedValue(fun.arg(2))
      c0 = E.constant_index(0)
      c32 = E.constant_index(32)
      with E.LoopNestContext([c0, c0, c0], [c32, c32, c32], [1, 1, 1]) as (i, j,
                                                                           k):
        C.store([i, j], A.load([i, k]) * B.load([k, j]))
      E.ret([])
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testMatrixMultiply
    #       CHECK: "affine.for"()
    #       CHECK: "affine.for"()
    #       CHECK: "affine.for"()
    #   CHECK-DAG:  %{{.*}} = "affine.load"
    #   CHECK-DAG:  %{{.*}} = "affine.load"
    #       CHECK:  %{{.*}} = mulf %{{.*}}, %{{.*}} : f32
    #       CHECK:  "affine.store"
    #  CHECK-SAME:  memref<32x32xf32>
    #       CHECK: {lower_bound = () -> (0), step = 1 : index, upper_bound = () -> (32)} : () -> ()
    #       CHECK: {lower_bound = () -> (0), step = 1 : index, upper_bound = () -> (32)} : () -> ()
    #       CHECK: {lower_bound = () -> (0), step = 1 : index, upper_bound = () -> (32)} : () -> ()

  def testRet(self):
    self.setUp()
    with self.module.function_context("foo", [],
                                      [self.indexType, self.indexType]) as fun:
      c42 = E.constant_index(42)
      c0 = E.constant_index(0)
      E.ret([c42, c0])
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testRet
    #       CHECK:    %{{.*}} = constant 42 : index
    #       CHECK:    %{{.*}} = constant 0 : index
    #       CHECK:    return %{{.*}}, %{{.*}} : index, index

  def testSelectOp(self):
    self.setUp()
    with self.module.function_context("foo", [self.boolType],
                                      [self.i32Type]) as fun:
      a = E.constant_int(42, 32)
      b = E.constant_int(0, 32)
      E.ret([E.select(fun.arg(0), a, b)])
      printWithCurrentFunctionName(str(fun))
    # CHECK-LABEL: testSelectOp
    #       CHECK:  %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : i32


# Until python 3.6 this cannot be used because the order in the dict is not the
# order of method declaration.
def runTests():
  def isTest(attr):
    return inspect.ismethod(attr) and "EdscTest.setUp " not in str(attr)

  edscTest = EdscTest()
  tests = sorted(filter(isTest,
                        (getattr(edscTest, attr) for attr in dir(edscTest))),
                 key = lambda x : str(x))
  for test in tests:
    test()

if __name__ == '__main__':
  runTests()
