<commit_before>package opencl.generator

import ir._
import ir.ast._
import lift.arithmetic.SizeVar
import opencl.executor._
import opencl.ir._
import opencl.ir.pattern._
import org.junit.Assert._
import org.junit.Assume.assumeFalse
import org.junit.{Ignore, Test}

object TestBarrier extends TestWithExecutor

class TestBarrier {
  @Test def basicBarrier(): Unit = {
    val inputSize = 1024
    val input = Array.tabulate(inputSize)(_.toFloat)

    // Barrier should be removed
    val f = fun(
      ArrayTypeWSWC(Float, SizeVar("N")),
      input => Join() o MapWrg(MapLcl(id)) o Split(128) $ input
    )

    val kernel = Compile(f)
    val (output: Array[Float], _) = Execute(inputSize)(kernel, f, input)

    assertArrayEquals(input, output, 0.0f)<commit_msg>Add some barrier testcases similar to the spmv problem<commit_after>package opencl.generator

import ir._
import ir.ast._
import lift.arithmetic.SizeVar
import opencl.executor._
import opencl.ir._
import opencl.ir.pattern._
import org.junit.Assert._
import org.junit.Assume.assumeFalse
import org.junit.{Ignore, Test}

object TestBarrier extends TestWithExecutor

class TestBarrier {

  @Ignore
  @Test
  def sequentialMapFollowsParallelWithLocalMemory(): Unit = {
    val inputSize = 1024
    val input = Array.tabulate(inputSize)(_.toFloat)

    // TODO: Can't have a barrier after MapSeq
    val f = fun(
      ArrayTypeWSWC(Float, SizeVar("N")),
      input => Join() o MapWrg(toGlobal(MapSeq(id)) o toLocal(MapLcl(id))) o Split(128) $ input
    )

    val kernel = Compile(f)
    val (output: Array[Float], _) = Execute(128)(kernel, f, input)

    assertArrayEquals(input, output, 0.0f)
    assertEquals(2, "barrier".r.findAllMatchIn(kernel).length)
  }

  @Test
  def sequentialMapFollowsParallel(): Unit = {
    val inputSize = 1024
    val input = Array.tabulate(inputSize)(_.toFloat)

    val f = fun(
      ArrayTypeWSWC(Float, SizeVar("N")),
      input => Join() o MapWrg(MapSeq(id) o MapLcl(id)) o Split(128) $ input
    )

    val kernel = Compile(f)
    val (output: Array[Float], _) = Execute(128)(kernel, f, input)

    assertArrayEquals(input, output, 0.0f)
    assertEquals(1, "barrier".r.findAllMatchIn(kernel).length)
  }

  @Ignore
  @Test
  def sequentialReduceFollowsParallelWithLocalMemory(): Unit = {
    val inputSize = 1024
    val input = Array.tabulate(inputSize)(_.toFloat)

    // TODO: Can't have a barrier after MapSeq or ReduceSeq
    val f = fun(
      ArrayTypeWSWC(Float, SizeVar("N")),
      input => Join() o MapWrg(toGlobal(MapSeq(id)) o ReduceSeq(add, 0.0f) o toLocal(MapLcl(id))) o Split(128) $ input
    )

    val kernel = Compile(f)
    val (output: Array[Float], _) = Execute(128)(kernel, f, input)

    val gold = input.grouped(128).map(_.sum).toArray

    assertArrayEquals(gold, output, 0.001f)
    assertEquals(2, "barrier".r.findAllMatchIn(kernel).length)
  }

  @Test
  def sequentialReduceFollowsParallel(): Unit = {
    val inputSize = 1024
    val input = Array.tabulate(inputSize)(_.toFloat)

    val f = fun(
      ArrayTypeWSWC(Float, SizeVar("N")),
      input => Join() o MapWrg(toGlobal(MapSeq(id)) o ReduceSeq(add, 0.0f) o MapLcl(id)) o Split(128) $ input
    )

    val kernel = Compile(f)
    val (output: Array[Float], _) = Execute(128)(kernel, f, input)

    val gold = input.grouped(128).map(_.sum).toArray

    assertArrayEquals(gold, output, 0.001f)
    assertEquals(1, "barrier".r.findAllMatchIn(kernel).length)
  }

  @Test def basicBarrier(): Unit = {
    val inputSize = 1024
    val input = Array.tabulate(inputSize)(_.toFloat)

    // Barrier should be removed
    val f = fun(
      ArrayTypeWSWC(Float, SizeVar("N")),
      input => Join() o MapWrg(MapLcl(id)) o Split(128) $ input
    )

    val kernel = Compile(f)
    val (output: Array[Float], _) = Execute(inputSize)(kernel, f, input)

    assertArrayEquals(input, output, 0.0f)
{'fraction_non_alphanumeric': 0.12896569512509673, 'fraction_numerical': 0.02166623678101625, 'mean_word_length': 3.672289156626506, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2483729', 'n_tokens_mistral': 1399, 'n_tokens_neox': 1302, 'n_words': 319}