];
}

def PrepareQuantizePass : Pass<"stablehlo-prepare-quantize", "mlir::ModuleOp"> {
  let summary = "Prepare StableHLO dialect for static range quantization by converting quantfork.stats into quantfork.qcast and dcast ops.";
  let options = [
    Option<"enable_per_channel_quantized_weight_",
        "enable-per-channel-quantized-weight",
        "bool", /*default=*/"true",
        "Whether to enable per-channel quantized weights.">,

      %1 = "quantfork.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
      %2 = "quantfork.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
      %3 = stablehlo.reshape %2 : (tensor<1x3xf32>) -> tensor<3x1xf32>

  %0 = "quantfork.qcast"(%cst_1) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.003937007874015748,0.003937007874015748}>>
  %1 = "quantfork.dcast"(%0) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32:3, {0.003937007874015748,0.003937007874015748}>>) -> tensor<2x3x3x2xf32>

		tpp.Param.goString(indent+2, "Param: "))
}

// Cast is a type cast.
type Cast struct {
	To AST
}

func (c *Cast) print(ps *printState) {
	ps.writeString("operator ")
	ps.print(c.To)
}

func (c *Cast) Traverse(fn func(AST) bool) {
	if fn(c) {
		c.To.Traverse(fn)
	}
}

func (c *Cast) Copy(fn func(AST) AST, skip func(AST) bool) AST {
	if skip(c) {

  // CHECK-NOT: %[[q0:.*]] = "quantfork.qcast"(%[[cst0]])
  // CHECK-NOT: %[[q1:.*]] = "quantfork.qcast"(%[[cst1]])
  // CHECK: %[[q2:.*]] = "quantfork.qcast"(%[[cst2]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.023529411764705882:-128>
  // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.023529411764705882:-128>
  // CHECK: %[[q3:.*]] = "quantfork.qcast"(%[[cst3]])

   %1 = "quantfork.qcast"(%0) {volatile} : (tensor<2x3xf32>) -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>
   %2 = "quantfork.dcast"(%1) : (tensor<2x3x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>) -> tensor<2x3xf32>
   %3 = "quantfork.qcast"(%arg0) {volatile} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 6.000000e-03:-128>>

  // CHECK-NOT: "quantfork.qcast"
  // CHECK-NOT: "quantfork.dcast"
  // CHECK: return %[[CST]]
  %cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32>
  %q = "quantfork.qcast"(%cst) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8:f32, 0.013075299590241675:-64>>

        rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(
            op, qtype, quantized_attr);
        return success();
      }
    }
    return failure();
  }
};

// Replaces quantfork.dcast with stablehlo.uniform_dequantize.
class ConvertDequantizeCastToUniformDequantizePattern
    : public OpRewritePattern<quantfork::DequantizeCastOp> {
 public:

    b.setInsertionPoint(op);
    if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
      auto dcast = b.create<quantfork::DequantizeCastOp>(
          dq.getLoc(), dq.getOutput().getType(), dq.getInput());
      dq.getOutput().replaceAllUsesWith(dcast);
      dq.erase();
    } else if (auto q = llvm::dyn_cast<QuantizeOp>(op)) {
      auto qcast = b.create<quantfork::QuantizeCastOp>(
          q.getLoc(), q.getOutput().getType(), q.getInput());

	"SAHF",
	"SALB",
	"SALL",
	"SALQ",
	"SALW",
	"SARB",
	"SARL",
	"SARQ",
	"SARW",
	"SARXL",
	"SARXQ",
	"SBBB",
	"SBBL",
	"SBBQ",
	"SBBW",
	"SCASB",
	"SCASL",
	"SCASQ",
	"SCASW",
	"SETCC",
	"SETCS",
	"SETEQ",
	"SETGE",
	"SETGT",
	"SETHI",
	"SETLE",
	"SETLS",
	"SETLT",
	"SETMI",
	"SETNE",
	"SETOC",
	"SETOS",
	"SETPC",