// RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s

gpu.module @test_elementwise_ops {

  // CHECK-LABEL: unary_ops_sg_layout_only
  gpu.func @unary_ops_sg_layout_only(%a: memref<24x32xf32>) {
    %tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    // CHECK: math.exp {{.*}} : vector<12x8xf32>
    %exp = math.exp %load_a
      : vector<24x32xf32>
    // CHECK: arith.negf {{.*}} : vector<12x8xf32>
    %negf = arith.negf %exp
      : vector<24x32xf32>
    %anchor = xegpu.convert_layout %negf
      <{
        input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>,
        target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>
      }> : vector<24x32xf32>
    gpu.return
  }

  // CHECK-LABEL: unary_ops
  gpu.func @unary_ops(%a: memref<24x32xf32>) {
    %tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    // CHECK: math.exp {{.*}} : vector<12x8xf32>
    %exp = math.exp %load_a
      : vector<24x32xf32>
    // CHECK: arith.negf {{.*}} : vector<12x8xf32>
    %negf = arith.negf %exp
      : vector<24x32xf32>
    %anchor = xegpu.convert_layout %negf
      <{
        input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
        target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
      }> : vector<24x32xf32>
    gpu.return
  }

  // CHECK-LABEL: binary_ops
  gpu.func @binary_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {
    %tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32> 
    %tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    // CHECK: arith.addf {{.*}}, {{.*}} : vector<12x8xf32>
    %addf = arith.addf %load_a, %load_b : vector<24x32xf32>
    // CHECK: math.powf {{.*}}, {{.*}} : vector<12x8xf32>
    %powf = math.powf %addf, %load_b : vector<24x32xf32>
    %anchor = xegpu.convert_layout %powf
      <{
        input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
        target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
      }> : vector<24x32xf32>
    gpu.return
  }

  // CHECK-LABEL: ternary_ops
  gpu.func @ternary_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>, %c: memref<24x32xi1>) {
    %tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %tdesc_c = xegpu.create_nd_tdesc %c : memref<24x32xi1> -> !xegpu.tensor_desc<24x32xi1>
    %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    %load_c = xegpu.load_nd %tdesc_c[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xi1> -> vector<24x32xi1>
    // CHECK: arith.select {{.*}}, {{.*}}, {{.*}} : vector<12x8xi1>, vector<12x8xf32>
    %select = arith.select %load_c, %load_a, %load_b
      : vector<24x32xi1>, vector<24x32xf32>
    // CHECK: math.fma  {{.*}}, {{.*}}, {{.*}} : vector<12x8xf32>
    %fma = math.fma %load_a, %load_b, %select
      : vector<24x32xf32>
    %anchor = xegpu.convert_layout %fma
      <{
        input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
        target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
      }> : vector<24x32xf32>
    gpu.return
  }

  // CHECK-LABEL: type_conversion_ops
  gpu.func @type_conversion_ops(%a: memref<24x32xf32>, %b: memref<24x32xi32>) {
    %tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
    %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
    // CHECK: arith.truncf {{.*}} : vector<12x8xf32> to vector<12x8xf16>
    %truncf = arith.truncf %load_a
      : vector<24x32xf32> to vector<24x32xf16>
    // CHECK: arith.bitcast {{.*}} : vector<12x8xf16> to vector<12x8xi16>
    %bitcast = arith.bitcast %truncf
      : vector<24x32xf16> to vector<24x32xi16>
    %anchor = xegpu.convert_layout %bitcast
      <{
        input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
        target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
      }> : vector<24x32xi16>
    gpu.return
  }

  // CHECK-LABEL: comparison_ops
  gpu.func @comparison_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>, %c: memref<24x32xi32>, %d: memref<24x32xi32>) {
    %tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %tdesc_c = xegpu.create_nd_tdesc %c : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
    %tdesc_d = xegpu.create_nd_tdesc %d : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
    %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    %load_c = xegpu.load_nd %tdesc_c[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
    %load_d = xegpu.load_nd %tdesc_d[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
    // CHECK: arith.cmpf ult, {{.*}}, {{.*}} : vector<12x8xf32>
    %cmpf = arith.cmpf ult, %load_a, %load_b
      : vector<24x32xf32>
    // CHECK: arith.cmpi eq, {{.*}}, {{.*}} : vector<12x8xi32>
    %cmpi = arith.cmpi eq, %load_c, %load_d
      : vector<24x32xi32>
    %res = arith.select %cmpi, %cmpi, %cmpf : vector<24x32xi1>, vector<24x32xi1>
    %anchor = xegpu.convert_layout %res
      <{
        input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
        target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
      }> : vector<24x32xi1>
    gpu.return
  }

  // 1 to N decomposition of elementwise operations
  // CHECK-LABEL: elementwise_ops_rr_assignment
  gpu.func @elementwise_ops_rr_assignment(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {
     %tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
    %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
      : !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
    // CHECK-COUNT-12: arith.negf {{.*}} : vector<2x2xf32>
    // CHECK-NOT: arith.negf
    %negf = arith.negf %load_a
      : vector<24x32xf32>
    // CHECK-COUNT-12: math.powf {{.*}}, {{.*}} : vector<2x2xf32>
    // CHECK-NOT: math.powf
    %powf = math.powf %negf, %load_b
      : vector<24x32xf32>
    %anchor = xegpu.convert_layout %powf
      <{
        input_layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>,
        target_layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>
      }> : vector<24x32xf32>
    gpu.return
  }
}